diff --git a/src/streaming/env.py b/src/streaming/env.py index f4c863a..79b6a9f 100644 --- a/src/streaming/env.py +++ b/src/streaming/env.py @@ -50,15 +50,34 @@ from utils.tool import logger +ASR_TIMEOUT_S = 30.0 # a 15 s segment normally transcribes in ~1-3 s +ASR_MAX_ATTEMPTS = 3 + + def transcribe_audio_segment(segment: AudioSegment, audio_format: str = "mp3") -> str: - """Transcribe a pydub ``AudioSegment`` using OpenAI Whisper (``whisper-1``).""" - buf = BytesIO() - segment.export(buf, format=audio_format) - buf.seek(0) - buf.name = f"audio.{audio_format}" - client = OpenAI() - transcript = client.audio.transcriptions.create(model="whisper-1", file=buf, language="en") - return (transcript.text or "").strip() + """Transcribe a pydub ``AudioSegment`` using OpenAI Whisper (``whisper-1``). + + The request carries an explicit timeout and is retried: without one, a hung + HTTP call blocks the listener thread for the rest of the turn (observed once + in run 138: the listener froze at 135 s of a 236 s speech and the main + thread then waited the full listener-join timeout). + """ + raw = BytesIO() + segment.export(raw, format=audio_format) + data = raw.getvalue() + client = OpenAI(timeout=ASR_TIMEOUT_S, max_retries=0) + last_err: Optional[Exception] = None + for attempt in range(ASR_MAX_ATTEMPTS): + buf = BytesIO(data) + buf.name = f"audio.{audio_format}" + try: + transcript = client.audio.transcriptions.create(model="whisper-1", file=buf, language="en") + return (transcript.text or "").strip() + except Exception as e: # timeout / transient API error + last_err = e + logger.warning(f"[StreamingInputEnv] ASR attempt {attempt + 1}/{ASR_MAX_ATTEMPTS} failed: {e}") + time.sleep(1.0 * (attempt + 1)) + raise RuntimeError(f"ASR failed after {ASR_MAX_ATTEMPTS} attempts: {last_err}") def _log_id_from_filename(filename: str) -> str: @@ -229,12 +248,56 @@ def run(self) -> None: except KeyboardInterrupt: logger.info("[StreamingInputEnv] KeyboardInterrupt; stopping.") finally: + # Drain: the speech is over, so transcribe whatever audio is left + # even if it is shorter than min_audio_seconds. Without this the + # final None: + """Transcribe the un-emitted tail of the continuous audio (any length >= 1 s).""" + cfg = self.config + log_id = "continuous" + available_audio = self._read_audio_up_to_cursor(continuous_file, cursor_seconds) + if available_audio is None: + return + already_ms = emitted_ms.get(log_id, 0) + end_ms = len(available_audio) + if cap_audio_ms is not None: + end_ms = min(end_ms, cap_audio_ms) + tail_ms = end_ms - already_ms + if tail_ms < 1000: + logger.debug(f"[StreamingInputEnv] drain_skip tail={tail_ms/1000.0:.2f}s t={time.time():.3f}") + return + segment = available_audio[already_ms:end_ms] + logger.debug(f"[StreamingInputEnv] drain_start audio_range={already_ms/1000.0:.2f}-{end_ms/1000.0:.2f}s t={time.time():.3f}") + asr_start = time.time() + text = transcribe_audio_segment(segment, audio_format=cfg.audio_format) + emitted_ms[log_id] = end_ms + logger.info( + f"[StreamingInputEnv] Drain transcribed s={tail_ms/1000.0:.2f} words={len(text.split())} " + f"total_processed={end_ms/1000.0:.2f}s asr_time={time.time()-asr_start:.2f}s" + ) + # force_flush=True: the speech is over, push the text into the tree now + self._append_transcript_text(log_id, text, True, text_buf, text_first_seen) + def _idle_text_flush(self, text_buf: Dict[str, List[str]], text_first_seen: Dict[str, float]) -> None: cfg = self.config if cfg.max_text_wait_seconds is None or cfg.max_text_wait_seconds <= 0: @@ -375,10 +438,22 @@ def _append_transcript_text( text_first_seen.pop(log_id, None) def _read_audio_up_to_cursor(self, continuous_file: Path, cursor_seconds: float) -> Optional[AudioSegment]: - """Read continuous audio file up to the cursor position.""" + """Read continuous audio file up to the cursor position. + + The decoded file is cached by (mtime, size): the playback thread only + rewrites it once per chunk, but this is polled every second and decoding + a 4-minute MP3 costs ~0.5 s of CPU each time. + """ try: read_start = time.time() - full_audio = AudioSegment.from_file(str(continuous_file)) + st = continuous_file.stat() + key = (st.st_mtime_ns, st.st_size) + cached = getattr(self, "_continuous_cache", None) + if cached is not None and cached[0] == key: + full_audio = cached[1] + else: + full_audio = AudioSegment.from_file(str(continuous_file)) + self._continuous_cache = (key, full_audio) cursor_ms = int(cursor_seconds * 1000) if cursor_ms <= 0: return None diff --git a/src/streaming/overlap.py b/src/streaming/overlap.py index bef2073..ffff63a 100644 --- a/src/streaming/overlap.py +++ b/src/streaming/overlap.py @@ -19,6 +19,7 @@ from __future__ import annotations import argparse +import os import json import sys import threading @@ -153,7 +154,11 @@ def _playback_main_loop( try: write_start = time.time() - continuous_audio.export(str(continuous_file), format=ext) + # Export to a temp file and atomically rename so the listener + # never opens a half-written file (ffmpeg decode errors). + tmp_file = continuous_file.with_name(continuous_file.name + ".tmp") + continuous_audio.export(str(tmp_file), format=ext) + os.replace(tmp_file, continuous_file) write_end = time.time() logger.debug( f"[PlaybackMain] file_write stage={stage} side={side} chunk_idx={next_idx} " @@ -169,11 +174,18 @@ def _playback_main_loop( f"duration={chunk_duration:.2f}s t={playback_start:.3f}" ) - # Simulate speaking in fixed increments + # Simulate speaking in fixed increments. The decode + continuous-file + # export above took real time (~1 s, growing with speech length); a + # real player would have started this chunk the moment its file was + # ready, so charge that bookkeeping time against the first sleep(s) + # instead of inserting a silent gap between every pair of chunks. elapsed_in_chunk = 0.0 + bookkeeping_s = max(0.0, playback_start - wait_end) while elapsed_in_chunk < chunk_duration: sleep_time = min(self._min_playback_increment, chunk_duration - elapsed_in_chunk) - time.sleep(sleep_time) + actual_sleep = max(0.0, sleep_time - bookkeeping_s) + bookkeeping_s = max(0.0, bookkeeping_s - sleep_time) + time.sleep(actual_sleep) # Advance cursor - listener can now access this audio playback_cursor[0] += sleep_time diff --git a/src/tts_streaming.py b/src/tts_streaming.py index 71e3a26..f144d26 100644 --- a/src/tts_streaming.py +++ b/src/tts_streaming.py @@ -13,6 +13,7 @@ import concurrent.futures import csv import json +import os import threading import time from dataclasses import dataclass, asdict @@ -27,7 +28,7 @@ from utils.tool import remove_citation, remove_subtitles from utils.time_estimator import LengthEstimator -from utils.fs_wrapper import FastSpeechWrapper +from utils.fs_wrapper import get_shared_wrapper # -------- config -------- TOLERANCE_RATIO = 0.10 # tolerance = target_s * this ratio (both sides) @@ -41,6 +42,27 @@ ABS_PRESTART_CHARS = 1000 # also pre-start when chunk[i+2] is absolutely large (>= this many chars), regardless of ratio SPEED_ADJUST_MIN = 0.85 # TTS speed clamp lower bound SPEED_ADJUST_MAX = 1.15 # TTS speed clamp upper bound +REFINE_DEADLINE_MARGIN_S = 2.0 # stop refining this much before the previous chunk finishes playing (decode/write/bridge/poll latency) +SPEED_ADJUST_MIN_SLACK_S = 4.0 # only re-synthesize with adjusted speed when at least this much playback slack remains +MAX_CHUNK_CHARS = 900 # paragraphs longer than this are split at sentence boundaries into ~equal pieces +TARGET_CHUNK_S = 40.0 # desired audio length per chunk when the text has too few paragraphs to stream +MIN_STREAM_CHUNKS = 3 # texts with fewer paragraphs than this are re-chunked by sentence packing +MAX_STREAM_CHUNKS = 8 +SEAM_HEAD_MS = 60 # leading silence kept at the start of every chunk +SEAM_TAIL_MS = 250 # trailing silence normalized to exactly this at the end of every chunk +SEAM_FADE_MS = 10 # tiny fade in/out so the cut never clicks +# Online calibration of FastSpeech duration estimates against the real TTS voice. +# FastSpeech (LJSpeech) and tts-1 "echo" differ in speaking rate by a text-dependent +# few percent; the static "*1.11-7" fit leaves ~5.5% mean abs error (p90 11.6%), +# which is the main reason a refined chunk still lands outside its +-10% tolerance. +# Each chunk gives one (estimate, actual) observation; an EMA of the ratio is +# applied to subsequent estimates. +# OFF by default: an offline replay on 206 chunks (runs 130-138) showed the +# residual is text-dependent noise with ~0 lag-1 autocorrelation, so chasing it +# with an EMA slightly *increases* error. Kept as an ablation (TTS_STREAM_CALIB=1). +CALIBRATION_ENABLED = os.environ.get("TTS_STREAM_CALIB", "0") == "1" +CALIB_EMA_ALPHA = 0.5 +CALIB_MIN, CALIB_MAX = 0.80, 1.25 # -------- dataclasses -------- @@ -107,8 +129,37 @@ def _in_range(est: float, target_s: float, tol_s: float, tol_upper_s: float) -> return (est - target_s) <= tol_upper_s and (target_s - est) <= tol_s +class _DurationCalibrator: + """Process-wide EMA of actual_tts_seconds / fastspeech_estimate_seconds.""" + + def __init__(self) -> None: + self.ratio = 1.0 + self.n = 0 + self._lock = threading.Lock() + + def observe(self, est_s: float, actual_s: float) -> None: + if not CALIBRATION_ENABLED or est_s <= 0 or actual_s <= 0: + return + r = actual_s / est_s + with self._lock: + self.ratio = r if self.n == 0 else (1 - CALIB_EMA_ALPHA) * self.ratio + CALIB_EMA_ALPHA * r + self.ratio = max(CALIB_MIN, min(CALIB_MAX, self.ratio)) + self.n += 1 + + def apply(self, est_s: float) -> float: + if not CALIBRATION_ENABLED: + return est_s + with self._lock: + return est_s * self.ratio + + +_CALIB = _DurationCalibrator() + + def _fastspeech_estimate(text: str) -> float: - wrapper = FastSpeechWrapper(batch_size=2) + # Shared, lazily-loaded model: constructing a wrapper per call reloads the + # checkpoint (+ vocoder) from disk every time (~5 s and GBs of GPU memory). + wrapper = get_shared_wrapper(batch_size=2) lengths = wrapper.query_time(text) length = lengths[0] if isinstance(lengths, list) else float(lengths) length = length * 1.11 - 7 if length > 100 else length @@ -170,6 +221,10 @@ def _revise_to_n_words(client, text: str, n_words: int, prev_texts: List[str], n }, ], max_completion_tokens=4096, + # A length-constrained rewrite needs no chain-of-thought: with the default + # effort gpt-5-mini spends ~700-1600 reasoning tokens (~10 s) per call; + # "minimal" gives the same word-count accuracy in ~2 s. + reasoning_effort="minimal", ) return (resp.choices[0].message.content or "").strip() @@ -216,10 +271,11 @@ def _tts_with_retry(client, content: str, voice: str = "echo", speed: float = 1. class _TtsCandidate: iteration: int # global index in shared candidate pool text: str - fs_estimated_s: float + fs_estimated_s: float # calibrated estimate used for the in-range decision future: Any # concurrent.futures.Future -> Dict from _query_time_profiled worker_label: str = "" # "prestart" / "normal" — which worker produced this intra_iter: int = 0 # iteration index within the producing worker (0 = raw, k = k-th refine) + fs_raw_s: float = 0.0 # uncalibrated FastSpeech estimate (for calibration updates) def _pick_best_completed( @@ -350,7 +406,7 @@ def aggregate_stats(self) -> Tuple[int, List[float], List[float]]: fs_times.extend(s["fs_times"]) return n_ref_total, llm_times, fs_times - def add_candidate(self, text: str, est: float, label: str, intra_iter: int) -> _TtsCandidate: + def add_candidate(self, text: str, est: float, label: str, intra_iter: int, est_raw: float = 0.0) -> _TtsCandidate: with self.candidates_lock: iteration = len(self.candidates) cand = _TtsCandidate( @@ -360,6 +416,7 @@ def add_candidate(self, text: str, est: float, label: str, intra_iter: int) -> _ future=self.executor.submit(_tts_with_retry, self.client, text, self.voice), worker_label=label, intra_iter=intra_iter, + fs_raw_s=est_raw, ) self.candidates.append(cand) return cand @@ -390,16 +447,17 @@ def _refine_worker(ctx: _ChunkRefineContext, label: str) -> None: # ---- step 0: fs estimate raw text + submit raw TTS candidate ---- t = _now() try: - est = _fastspeech_estimate(cur) + est_raw = _fastspeech_estimate(cur) except Exception: return + est = _CALIB.apply(est_raw) ctx.add_fs_time(label, _now() - t) if ctx.stop_event.is_set(): return target_s, tol_s, tol_upper_s = ctx.get_target() - cand = ctx.add_candidate(cur, est, label, intra_iter=0) + cand = ctx.add_candidate(cur, est, label, intra_iter=0, est_raw=est_raw) if _in_range(est, target_s, tol_s, tol_upper_s): try: @@ -432,15 +490,16 @@ def _refine_worker(ctx: _ChunkRefineContext, label: str) -> None: t = _now() try: - est = _fastspeech_estimate(cur) + est_raw = _fastspeech_estimate(cur) except Exception: break + est = _CALIB.apply(est_raw) ctx.add_fs_time(label, _now() - t) if ctx.stop_event.is_set(): break - cand = ctx.add_candidate(cur, est, label, intra_iter=n_ref_local) + cand = ctx.add_candidate(cur, est, label, intra_iter=n_ref_local, est_raw=est_raw) target_s, tol_s, tol_upper_s = ctx.get_target() if _in_range(est, target_s, tol_s, tol_upper_s): @@ -453,6 +512,28 @@ def _refine_worker(ctx: _ChunkRefineContext, label: str) -> None: # -------- chunk utilities -------- +def _normalize_seam_silence(seg: AudioSegment) -> Tuple[AudioSegment, bool]: + """Trim/pad head and tail silence so every chunk boundary sounds the same. + + tts-1 output carries 0-150 ms of leading and 100-1200 ms of trailing silence; + concatenated chunk-by-chunk that yields pauses of 0.2-1.3 s at the seams. + Returns (segment, changed). + """ + from pydub.silence import detect_leading_silence + + if len(seg) < 1000: + return seg, False + lead = detect_leading_silence(seg, silence_threshold=-40, chunk_size=5) + trail = detect_leading_silence(seg.reverse(), silence_threshold=-40, chunk_size=5) + if len(seg) - lead - trail < 500: # no speech detected (e.g. silent fallback) -> leave it + return seg, False + start = max(0, lead - SEAM_HEAD_MS) + end = len(seg) - trail + core = seg[start:end] + out = core.fade_in(SEAM_FADE_MS).fade_out(SEAM_FADE_MS) + AudioSegment.silent(duration=SEAM_TAIL_MS, frame_rate=seg.frame_rate) + return out, True + + def _split_sentences(text: str) -> List[str]: """Split text into sentences on '.', '!', '?' boundaries.""" import re @@ -515,6 +596,68 @@ def split_by_paragraphs(text: str) -> List[str]: return parts if parts else [text.strip()] if text.strip() else [] +def _pack_sentences(text: str, target_chars: int) -> List[str]: + """Greedily pack whole sentences into pieces of roughly ``target_chars`` characters. + + Never cuts inside a sentence; a single sentence longer than the target becomes + its own piece. Pieces are balanced so the last one is not a tiny remainder. + """ + sentences = _split_sentences(text) + if not sentences: + return [] + total = sum(len(x) + 1 for x in sentences) + n_pieces = max(1, round(total / max(target_chars, 1))) + per_piece = total / n_pieces + pieces: List[str] = [] + cur: List[str] = [] + cur_len = 0 + for sent in sentences: + if cur and cur_len + len(sent) + 1 > per_piece * 1.15 and len(pieces) < n_pieces - 1: + pieces.append(" ".join(cur)) + cur, cur_len = [], 0 + cur.append(sent) + cur_len += len(sent) + 1 + if cur: + pieces.append(" ".join(cur)) + return pieces + + +def split_into_chunks(text: str, total_budget_s: float) -> List[str]: + """Split a speech into streamable chunks. + + Paragraphs are the preferred unit (they are natural rhetorical units and the + refinement prompt talks about "paragraphs"). Two failure modes of pure + paragraph splitting are handled here: + 1. The model emitted (almost) no paragraph breaks -> a single huge chunk, + which disables streaming and refinement entirely. In that case the text + is re-chunked by packing sentences to ~TARGET_CHUNK_S of audio each. + 2. One paragraph is much longer than the rest (> MAX_CHUNK_CHARS) -> it is + split at sentence boundaries into a few balanced pieces. + """ + paras = split_by_paragraphs(text) + if not paras: + return [] + total_chars = sum(len(p) for p in paras) + + if len(paras) < MIN_STREAM_CHUNKS: + n_target = int(round(total_budget_s / TARGET_CHUNK_S)) if total_budget_s > 0 else MIN_STREAM_CHUNKS + n_target = max(MIN_STREAM_CHUNKS, min(MAX_STREAM_CHUNKS, n_target)) + target_chars = max(200, total_chars // n_target) + out: List[str] = [] + for p in paras: + out.extend(_pack_sentences(p, target_chars)) + return out or paras + + out = [] + for p in paras: + if len(p) > MAX_CHUNK_CHARS: + n_pieces = -(-len(p) // MAX_CHUNK_CHARS) # ceil + out.extend(_pack_sentences(p, max(200, len(p) // n_pieces))) + else: + out.append(p) + return out + + # -------- pipeline -------- def run_pipeline( client, @@ -719,7 +862,10 @@ def _kickoff_last_chunk_prestart(iter_i: int) -> None: time_budget_s = target_s refined = chunk n_ref_used = 0 - fs_estimated_s = 0.0 + try: + fs_estimated_s = _fastspeech_estimate(chunk) # ~10 ms; seeds the calibrator below + except Exception: + fs_estimated_s = 0.0 refine_total_s = 0.0 in_range = True timed_out = False @@ -752,6 +898,9 @@ def _kickoff_last_chunk_prestart(iter_i: int) -> None: iter_tts_times_s = json.dumps([round(tts_api_s, 3)]) total_elapsed_s = tts_api_s overrun_s = tts_api_s + # chunk 0 is never refined, but record whether it actually hit its target + in_range = _in_range(audio_seconds, target_s, tol_s, tol_upper_s) + _CALIB.observe(fs_estimated_s, audio_seconds) # ---- chunks 1+: shared candidate pool with prestart + normal workers ---- else: @@ -790,7 +939,12 @@ def _kickoff_last_chunk_prestart(iter_i: int) -> None: normal_th.start() # Wait for ANY worker to find an ok candidate, OR until deadline. - ctx.done_event.wait(timeout=max(0.0, time_budget_s)) + # The deadline is the previous chunk's playback length minus a margin + # for everything that still has to happen after selection (decode, + # file write, bridge copy, playback-side stability poll); without the + # margin every timed-out chunk arrives a few seconds late and the + # listener hears a gap. + ctx.done_event.wait(timeout=max(0.0, time_budget_s - REFINE_DEADLINE_MARGIN_S)) ctx.stop_event.set() total_elapsed_s = _now() - ctx.t_start_wall @@ -812,16 +966,24 @@ def _kickoff_last_chunk_prestart(iter_i: int) -> None: in_range = False timed_out = True - # Speed adjustment if still out of range + # Speed adjustment if still out of range -- but only when there is + # enough playback slack left for another TTS round trip; otherwise + # the extra call itself causes an audible gap. target_now, tol_now, tol_upper_now = ctx.get_target() audio_s = float(tts_out["audio_seconds"]) - if not _in_range(audio_s, target_now, tol_now, tol_upper_now): + _CALIB.observe(chosen_cand.fs_raw_s, audio_s) + slack_s = time_budget_s - (_now() - chunk_t0) + if not _in_range(audio_s, target_now, tol_now, tol_upper_now) and slack_s >= SPEED_ADJUST_MIN_SLACK_S: raw_speed = audio_s / target_now if target_now > 0 else 1.0 clamped = max(SPEED_ADJUST_MIN, min(SPEED_ADJUST_MAX, raw_speed)) if abs(clamped - 1.0) > 0.01: try: speed_tts_out = _tts_with_retry(client, chosen_cand.text, voice=voice, speed=clamped) if abs(speed_tts_out["audio_seconds"] - target_now) < abs(audio_s - target_now): + print( + f" [speed-adjust] chunk {i}: speed={clamped:.2f} " + f"audio {audio_s:.1f}s -> {speed_tts_out['audio_seconds']:.1f}s (target {target_now:.1f}s)" + ) tts_out = speed_tts_out audio_s = float(tts_out["audio_seconds"]) except Exception: @@ -901,9 +1063,9 @@ def _worker_tts_in_order(label: str) -> List[float]: overrun_s = max(0.0, total_elapsed_s - time_budget_s) - # Best-effort cleanup (workers will exit at next stop_event check) - for w in ctx.workers: - w.join(timeout=2) + # Best-effort cleanup: workers are daemon threads that exit at their + # next stop_event check. Do NOT join them here -- a worker blocked in + # an LLM/TTS call would stall the main pipeline by the join timeout. ctx.executor.shutdown(wait=False) try: @@ -913,6 +1075,14 @@ def _worker_tts_in_order(label: str) -> List[float]: warnings.warn(f"Chunk {i}: pydub decode failed: {e}") seg = AudioSegment.silent(duration=int(audio_seconds * 1000)) + # ---- seam normalization: uniform head/tail silence on every chunk ---- + seg, seam_changed = _normalize_seam_silence(seg) + if seam_changed: + buf = BytesIO() + seg.export(buf, format="mp3") + mp3_bytes = buf.getvalue() + audio_seconds = len(seg) / 1000.0 + # ---- budget tracking ---- audio_budget_remaining -= audio_seconds + overrun_s prev_audio_s = audio_seconds @@ -921,7 +1091,10 @@ def _worker_tts_in_order(label: str) -> List[float]: if out_dir is not None: (out_dir / f"chunk_{i:03d}.txt").write_text(refined, encoding="utf-8") - (out_dir / f"chunk_{i:03d}.mp3").write_bytes(mp3_bytes) + # write atomically so the copy bridge never sees a partial file + tmp_mp3 = out_dir / f"chunk_{i:03d}.mp3.tmp" + tmp_mp3.write_bytes(mp3_bytes) + tmp_mp3.replace(out_dir / f"chunk_{i:03d}.mp3") combined_audio += seg @@ -988,7 +1161,7 @@ def _worker_tts_in_order(label: str) -> List[float]: f"fs_est={fs_estimated_s:.1f}s | actual={audio_seconds:.1f}s | " f"in_range={in_range} | timed_out={timed_out} | " f"cands={n_candidates_submitted}(done={n_candidates_done},used={used_candidate_iter}) | " - f"overrun={overrun_s:.2f}s | remaining={audio_budget_remaining:.1f}s" + f"overrun={overrun_s:.2f}s | remaining={audio_budget_remaining:.1f}s | calib={_CALIB.ratio:.3f}" ) i += 1 @@ -1055,7 +1228,7 @@ def convert_text_to_speech_streaming( audio_content, _ = remove_citation(content) audio_content = remove_subtitles(audio_content) - segments = split_by_paragraphs(audio_content) + segments = split_into_chunks(audio_content, total_budget_s) client = OpenAI() output_path = Path(output_path) diff --git a/src/utils/fs_wrapper.py b/src/utils/fs_wrapper.py index 99f55c0..446d7ce 100644 --- a/src/utils/fs_wrapper.py +++ b/src/utils/fs_wrapper.py @@ -1,5 +1,6 @@ import os import sys +import threading import numpy as np import torch @@ -8,13 +9,68 @@ fastspeech_path = os.path.join(os.path.dirname(__file__), "..", "..", "dependencies") sys.path.append(fastspeech_path) -from fastspeech2.synthesize import preprocess_english +import re +from string import punctuation + +import fastspeech2.transformer.Models as _fs_models +from fastspeech2.synthesize import preprocess_english, read_lexicon +from fastspeech2.text import text_to_sequence from fastspeech2.utils.model import get_model_2, get_vocoder -from fastspeech2.utils.tools import synth_samples, synth_samples_for_length, to_device +from fastspeech2.utils.tools import get_mask_from_lengths, synth_samples, synth_samples_for_length, to_device +from g2p_en import G2p root = f"{fastspeech_path}/fastspeech2/" +def _fast_sinusoid_encoding_table(n_position, d_hid, padding_idx=None): + """Vectorized drop-in for fastspeech2.transformer.Models.get_sinusoid_encoding_table. + + The upstream version evaluates n_position * d_hid Python calls (>1M for a + long chunk) every forward pass whose length exceeds max_seq_len; that alone + cost ~1 s per duration query. + """ + pos = np.arange(n_position, dtype=np.float64)[:, None] + hid = np.arange(d_hid, dtype=np.float64)[None, :] + table = pos / np.power(10000, 2 * (hid // 2) / d_hid) + table[:, 0::2] = np.sin(table[:, 0::2]) + table[:, 1::2] = np.cos(table[:, 1::2]) + if padding_idx is not None: + table[padding_idx] = 0.0 + return torch.FloatTensor(table) + + +_fs_models.get_sinusoid_encoding_table = _fast_sinusoid_encoding_table + +_LEXICON_CACHE = {} +_G2P = None +_PREP_LOCK = threading.Lock() + + +def _preprocess_english_cached(text, preprocess_config): + """Same phoneme sequence as fastspeech2.synthesize.preprocess_english, but the + 206k-line lexicon and the G2p model are loaded once instead of per call.""" + global _G2P + lex_path = preprocess_config["path"]["lexicon_path"] + with _PREP_LOCK: + if lex_path not in _LEXICON_CACHE: + _LEXICON_CACHE[lex_path] = read_lexicon(lex_path) + if _G2P is None: + _G2P = G2p() + lexicon = _LEXICON_CACHE[lex_path] + g2p = _G2P + text = text.rstrip(punctuation) + phones = [] + for w in re.split(r"([,;.\-\?\!\s+])", text): + if w.lower() in lexicon: + phones += lexicon[w.lower()] + else: + phones += list(filter(lambda p: p != " ", g2p(w))) + phones = "{" + "}{".join(phones) + "}" + phones = re.sub(r"\{[^\w\s]?\}", "{sp}", phones) + phones = phones.replace("}{", " ") + return np.array(text_to_sequence(phones, preprocess_config["preprocessing"]["text"]["text_cleaners"])) + + def pad_1D(inputs, PAD=0): def pad_data(x, length, PAD): x_padded = np.pad(x, (0, length - x.shape[0]), mode="constant", constant_values=PAD) @@ -27,7 +83,7 @@ def pad_data(x, length, PAD): class FastSpeechWrapper: - def __init__(self, batch_size=8): + def __init__(self, batch_size=8, load_vocoder=False): preprocess_config = yaml.load(open(f"{root}/config/LJSpeech/preprocess.yaml", "r"), Loader=yaml.FullLoader) model_config = yaml.load(open(f"{root}/config/LJSpeech/model.yaml", "r"), Loader=yaml.FullLoader) train_config = yaml.load(open(f"{root}/config/LJSpeech/train.yaml", "r"), Loader=yaml.FullLoader) @@ -39,8 +95,12 @@ def __init__(self, batch_size=8): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model = get_model_2(restore_step, self.configs, self.device, train=False) - self.vocoder = get_vocoder(model_config, self.device) + # The vocoder is only needed to synthesize waveforms. For duration + # estimation the mel length already determines the audio length exactly + # (len(wav) == mel_len * hop_length), so skip loading it by default. + self.vocoder = get_vocoder(model_config, self.device) if load_vocoder else None self.batch_size = batch_size + self._lock = threading.Lock() def process_text(self, texts): batchs = [] @@ -48,7 +108,7 @@ def process_text(self, texts): ids = f"Part-{i}" raw_text = text speakers = 0 - post_text = np.array(preprocess_english(text, self.configs[0])) + post_text = _preprocess_english_cached(text, self.configs[0]) text_len = len(post_text) batchs.append((ids, raw_text, speakers, post_text, text_len)) return batchs @@ -76,22 +136,60 @@ def synthesize(self, batchs): data = (ids, raw_texts, speakers, texts, text_lens, max(text_lens)) data = to_device(data, self.device) - with torch.no_grad(): - # Forward - output = self.model( - *(data[2:]), p_control=pitch_control, e_control=energy_control, d_control=duration_control - ) - length = synth_samples_for_length( - data, - output, - self.vocoder, - model_config, - preprocess_config, - train_config["path"]["result_path"], - ) + with self._lock, torch.no_grad(): + if self.vocoder is None: + # Duration-only path: encoder + variance adaptor give mel_len; + # the decoder/postnet/vocoder do not change it. Same value + # synth_samples_for_length() returns (wav len == mel_len * hop). + hop = preprocess_config["preprocessing"]["stft"]["hop_length"] + sr = preprocess_config["preprocessing"]["audio"]["sampling_rate"] + mel_lens = self._predict_mel_lens( + *(data[2:]), p_control=pitch_control, e_control=energy_control, d_control=duration_control + ) + length = [float(m) * hop / sr for m in mel_lens] + else: + output = self.model( + *(data[2:]), p_control=pitch_control, e_control=energy_control, d_control=duration_control + ) + length = synth_samples_for_length( + data, + output, + self.vocoder, + model_config, + preprocess_config, + train_config["path"]["result_path"], + ) lengths.extend(length) return lengths + def _predict_mel_lens(self, speakers, texts, src_lens, max_src_len, p_control=1.0, e_control=1.0, d_control=1.0): + """First half of FastSpeech2.forward (encoder + variance adaptor), returning predicted mel lengths.""" + m = self.model + src_masks = get_mask_from_lengths(src_lens, max_src_len) + output = m.encoder(texts, src_masks) + if m.speaker_emb is not None: + output = output + m.speaker_emb(speakers).unsqueeze(1).expand(-1, max_src_len, -1) + (_, _, _, _, _, mel_lens, _) = m.variance_adaptor( + output, src_masks, None, None, None, None, None, p_control, e_control, d_control + ) + return mel_lens.detach().cpu().numpy() + + +_SHARED_WRAPPER = None +_SHARED_WRAPPER_LOCK = threading.Lock() + + +def get_shared_wrapper(batch_size=8): + """Process-wide FastSpeechWrapper. Loading the checkpoint takes seconds and + ~GBs of GPU memory, so callers that only need duration estimates must share + one instance instead of constructing a new wrapper per query.""" + global _SHARED_WRAPPER + if _SHARED_WRAPPER is None: + with _SHARED_WRAPPER_LOCK: + if _SHARED_WRAPPER is None: + _SHARED_WRAPPER = FastSpeechWrapper(batch_size=batch_size) + return _SHARED_WRAPPER + if __name__ == "__main__": texts = [ diff --git a/src/utils/prompts/others.py b/src/utils/prompts/others.py index ec14a32..46159e1 100644 --- a/src/utils/prompts/others.py +++ b/src/utils/prompts/others.py @@ -15,9 +15,9 @@ ## Debate Format: 1. Two debaters or teams present arguments for and against a topic. 2. There are three stages: - * Opening statements (4 minutes, around 520 words) - * Rebuttal statements (4 minutes, around 520 words) - * Closing statements (2 minutes, around 260 words) + * Opening statements (4 minutes, around 580 words) + * Rebuttal statements (4 minutes, around 580 words) + * Closing statements (2 minutes, around 290 words) 3. Speaking order for each stage: FOR side speaks first, followed by AGAINST 4. Victory determination: * Audience votes before debate (initial position) diff --git a/src/utils/time_estimator.py b/src/utils/time_estimator.py index 2570158..eb0bf4b 100644 --- a/src/utils/time_estimator.py +++ b/src/utils/time_estimator.py @@ -10,7 +10,7 @@ from openai import OpenAI from .constants import openai_api_key -from .fs_wrapper import FastSpeechWrapper +from .fs_wrapper import get_shared_wrapper from .tool import remove_citation, remove_subtitles @@ -18,14 +18,14 @@ class LengthEstimator: def __init__(self, mode): self.mode = mode if self.mode == "fastspeech": - self.client = FastSpeechWrapper(batch_size=8) + self.client = get_shared_wrapper(batch_size=8) elif self.mode == "openai": self.client = OpenAI(api_key=openai_api_key) def query_time(self, content: List[str], mode=None) -> List[float]: if mode is not None and mode != self.mode: if self.mode == "fastspeech": - self.client = FastSpeechWrapper(batch_size=2) + self.client = get_shared_wrapper(batch_size=8) elif self.mode == "openai": self.client = OpenAI(api_key=openai_api_key)