OpenMOSS-Team/MOSS-TTS-Realtime
6
1import spaces2 3import argparse4import base645import functools6import json7import sys8import threading9import time10from collections import OrderedDict11from dataclasses import dataclass12from pathlib import Path13from typing import Iterator, Sequence14 15import gradio as gr16import numpy as np17 18import os19os.environ["TORCHDYNAMO_DISABLE"] = "1"20 21import torch22import torchaudio23import torch._dynamo24from transformers import AutoModel, AutoTokenizer25from mossttsrealtime import MossTTSRealtime, MossTTSRealtimeProcessor26from mossttsrealtime.streaming_mossttsrealtime import (27 AudioStreamDecoder,28 MossTTSRealtimeInference,29 MossTTSRealtimeStreamingSession,30)31 32torch._dynamo.config.cache_size_limit = 6433 34 35 36APP_DIR = Path(__file__).resolve().parent37AUDIO_DIR = APP_DIR / "asset"38LOG_DIR = APP_DIR / "logs"39SAMPLE_RATE = 2400040 41CODEC_MODEL_PATH = "OpenMOSS-Team/MOSS-Audio-Tokenizer"42MODEL_PATH = "OpenMOSS-Team/MOSS-TTS-Realtime"43TOKENIZER_PATH = "OpenMOSS-Team/MOSS-TTS-Realtime"44 45PROMPT_WAV = "asset/prompt_audio.mp3"46USER_WAV = "asset/user1.wav"47 48WARMUP_POLL_INTERVAL_SECONDS = 0.549DEFAULT_REPETITION_WINDOW = 5050WARMUP_STEP_TOKENS = DEFAULT_REPETITION_WINDOW + 151WARMUP_USER_TEXT = "Hello!"52WARMUP_BASE_ASSISTANT_TEXT = (53 "This startup warmup request primes the streaming text to speech path "54 "so the first real user request avoids the cold compile stall."55)56 57 58def _apply_seed(seed: int | None) -> None:59 if seed is None:60 return61 # ZeroGPU: avoid touching torch.cuda outside the managed GPU call.62 torch.manual_seed(seed)63 64 65def _load_audio(path: Path, target_sample_rate: int = SAMPLE_RATE) -> torch.Tensor:66 wav, sr = torchaudio.load(path)67 if sr != target_sample_rate:68 wav = torchaudio.functional.resample(wav, sr, target_sample_rate)69 if wav.shape[0] > 1:70 wav = wav.mean(dim=0, keepdim=True)71 return wav72 73 74def _load_codec(device: torch.device, codec_model_path: str):75 codec = AutoModel.from_pretrained(codec_model_path, trust_remote_code=True).eval()76 return codec.to(device)77 78 79def _extract_codes(encode_result):80 if isinstance(encode_result, dict):81 codes = encode_result["audio_codes"]82 83 elif isinstance(encode_result, (list, tuple)) and encode_result:84 codes = encode_result[0]85 else:86 codes = encode_result87 88 if isinstance(codes, np.ndarray):89 codes = torch.from_numpy(codes)90 91 if isinstance(codes, torch.Tensor) and codes.dim() == 3:92 if codes.shape[1] == 1:93 codes = codes[:, 0, :]94 elif codes.shape[0] == 1:95 codes = codes[0]96 else:97 raise ValueError(f"Unsupported 3D audio code shape: {tuple(codes.shape)}")98 99 return codes100 101 102@dataclass(frozen=True)103class BackendPaths:104 model_path: str105 tokenizer_path: str106 codec_model_path: str107 device_str: str108 attn_impl: str109 110 111@dataclass(frozen=True)112class GenerationConfig:113 temperature: float114 top_p: float115 top_k: int116 repetition_penalty: float117 repetition_window: int118 do_sample: bool119 max_length: int120 seed: int | None121 122 123@dataclass(frozen=True)124class StreamingConfig:125 text_chunk_tokens: int126 input_delay: float127 decode_chunk_frames: int128 decode_overlap_frames: int129 chunk_duration: float130 prebuffer_seconds: float131 buffer_threshold_seconds: float = 0.0132 133 134@dataclass(frozen=True)135class StreamingRequest:136 user_text: str137 assistant_text: str138 prompt_audio: str | None139 user_audio: str | None140 use_default_prompt: bool141 use_default_user: bool142 generation: GenerationConfig143 streaming: StreamingConfig144 backend: BackendPaths145 146 147@dataclass(frozen=True)148class StreamEvent:149 message: str150 audio: tuple[int, np.ndarray] | None = None151 152 153@dataclass(frozen=True)154class WarmupSnapshot:155 state: str156 progress: float157 message: str158 detail: str | None = None159 error: str | None = None160 161 @property162 def ready(self) -> bool:163 return self.state == "ready"164 165 @property166 def failed(self) -> bool:167 return self.state == "failed"168 169 170def _make_log_path(prefix: str) -> Path:171 LOG_DIR.mkdir(parents=True, exist_ok=True)172 stamp = time.strftime("%Y%m%d_%H%M%S", time.localtime())173 return LOG_DIR / f"{prefix}_{stamp}_{time.time_ns() % 1_000_000_000:09d}.jsonl"174 175 176def _compute_rtf_metrics(sample_count: int, sample_rate: int, started_at: float) -> dict[str, float | None]:177 elapsed_s = max(0.0, time.monotonic() - started_at)178 audio_s = float(sample_count) / float(sample_rate) if sample_count > 0 and sample_rate > 0 else 0.0179 rtf = (elapsed_s / audio_s) if audio_s > 0 else None180 return {181 "elapsed_s": elapsed_s,182 "audio_s": audio_s,183 "rtf": rtf,184 }185 186 187class StreamRTFLogger:188 def __init__(self, path: Path, started_at: float):189 self.path = path190 self.started_at = started_at191 self.chunk_count = 0192 self.sample_rate = SAMPLE_RATE193 self.samples_emitted = 0194 195 @classmethod196 def create(cls, request: "StreamingRequest", started_at: float) -> "StreamRTFLogger":197 logger = cls(_make_log_path("rtf"), started_at)198 logger.log_request_started(request)199 print(f"[MossTTSRealtime][rtf-log] {logger.path}", flush=True)200 return logger201 202 def log_request_started(self, request: "StreamingRequest") -> None:203 self._append(204 {205 "event": "request_started",206 "user_text_chars": len(request.user_text),207 "assistant_text_chars": len(request.assistant_text),208 "text_chunk_tokens": request.streaming.text_chunk_tokens,209 "decode_chunk_frames": request.streaming.decode_chunk_frames,210 "decode_overlap_frames": request.streaming.decode_overlap_frames,211 "chunk_duration_s": request.streaming.chunk_duration,212 "prebuffer_seconds": request.streaming.prebuffer_seconds,213 "temperature": request.generation.temperature,214 "top_p": request.generation.top_p,215 "top_k": request.generation.top_k,216 "repetition_penalty": request.generation.repetition_penalty,217 "repetition_window": request.generation.repetition_window,218 "do_sample": request.generation.do_sample,219 "max_length": request.generation.max_length,220 "seed": request.generation.seed,221 "device": request.backend.device_str,222 "attn_implementation": request.backend.attn_impl,223 }224 )225 226 def log_chunk(227 self,228 *,229 event_message: str,230 sample_rate: int,231 chunk: np.ndarray,232 first_audio_time: float | None,233 ) -> None:234 chunk = np.asarray(chunk).reshape(-1)235 if chunk.size == 0:236 return237 self.chunk_count += 1238 self.sample_rate = int(sample_rate)239 self.samples_emitted += int(chunk.size)240 metrics = _compute_rtf_metrics(self.samples_emitted, self.sample_rate, self.started_at)241 record = {242 "event": "stream_chunk",243 "message": event_message,244 "chunk_idx": self.chunk_count,245 "chunk_audio_s": float(chunk.size) / float(self.sample_rate),246 "audio_s_emitted": metrics["audio_s"],247 "elapsed_s": metrics["elapsed_s"],248 "rtf": metrics["rtf"],249 }250 if first_audio_time is not None:251 record["time_to_first_audio_ms"] = max(0.0, (first_audio_time - self.started_at) * 1000.0)252 self._append(record)253 254 def log_completion(self, *, first_audio_time: float | None) -> None:255 metrics = _compute_rtf_metrics(self.samples_emitted, self.sample_rate, self.started_at)256 record = {257 "event": "stream_complete",258 "chunk_count": self.chunk_count,259 "audio_s_total": metrics["audio_s"],260 "elapsed_s": metrics["elapsed_s"],261 "rtf": metrics["rtf"],262 }263 if first_audio_time is not None:264 record["time_to_first_audio_ms"] = max(0.0, (first_audio_time - self.started_at) * 1000.0)265 self._append(record)266 267 def log_no_audio(self) -> None:268 metrics = _compute_rtf_metrics(0, self.sample_rate, self.started_at)269 self._append(270 {271 "event": "stream_complete",272 "chunk_count": 0,273 "audio_s_total": 0.0,274 "elapsed_s": metrics["elapsed_s"],275 "rtf": None,276 "warning": "No audio chunks emitted.",277 }278 )279 280 def log_error(self, exc: Exception, *, first_audio_time: float | None) -> None:281 metrics = _compute_rtf_metrics(self.samples_emitted, self.sample_rate, self.started_at)282 record = {283 "event": "stream_error",284 "error_type": type(exc).__name__,285 "error": str(exc),286 "chunk_count": self.chunk_count,287 "audio_s_emitted": metrics["audio_s"],288 "elapsed_s": metrics["elapsed_s"],289 "rtf": metrics["rtf"],290 }291 if first_audio_time is not None:292 record["time_to_first_audio_ms"] = max(0.0, (first_audio_time - self.started_at) * 1000.0)293 self._append(record)294 295 def _append(self, payload: dict[str, object]) -> None:296 record = {297 "ts": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),298 **payload,299 }300 with self.path.open("a", encoding="utf-8") as handle:301 handle.write(json.dumps(record, ensure_ascii=False) + "\n")302 303 304class TokenChunkStream:305 def __init__(306 self,307 tokens: Sequence[int],308 chunk_size: int,309 ):310 self._tokens = list(tokens)311 self._chunk_size = int(chunk_size)312 313 def __iter__(self) -> Iterator[list[int]]:314 if not self._tokens:315 return316 step = len(self._tokens) if self._chunk_size <= 0 else self._chunk_size317 for idx in range(0, len(self._tokens), step):318 yield self._tokens[idx : idx + step]319 320 321class BufferedAudioTracker:322 def __init__(self, sample_rate: int):323 self.sample_rate = sample_rate324 self.start_time: float | None = None325 self.samples_emitted = 0326 327 def add_chunk(self, chunk: np.ndarray) -> None:328 if chunk.size == 0:329 return330 if self.start_time is None:331 self.start_time = time.monotonic()332 self.samples_emitted += int(chunk.size)333 334 def buffered_seconds(self) -> float:335 if self.start_time is None:336 return 0.0337 elapsed = time.monotonic() - self.start_time338 buffered = self.samples_emitted / self.sample_rate - elapsed339 return max(0.0, buffered)340 341 342class AudioFrameDecoder:343 def __init__(344 self,345 decoder: AudioStreamDecoder,346 codebook_size: int,347 audio_eos_token: int,348 ):349 self.decoder = decoder350 self.codebook_size = codebook_size351 self.audio_eos_token = audio_eos_token352 353 def decode_frames(self, audio_frames: list[torch.Tensor]) -> Iterator[np.ndarray]:354 for frame in audio_frames:355 tokens = frame356 if tokens.dim() == 3:357 tokens = tokens[0]358 if tokens.dim() != 2:359 raise ValueError(f"Expected [T, C] audio tokens, got {tuple(tokens.shape)}")360 tokens, stop = _sanitize_tokens(tokens, self.codebook_size, self.audio_eos_token)361 if tokens.numel() == 0:362 if stop:363 break364 continue365 self.decoder.push_tokens(tokens.detach())366 for wav in self.decoder.audio_chunks():367 if wav.numel() == 0:368 continue369 yield wav.detach().cpu().numpy().reshape(-1)370 if stop:371 break372 373 def flush(self) -> Iterator[np.ndarray]:374 final_chunk = self.decoder.flush()375 if final_chunk is not None and final_chunk.numel() > 0:376 yield final_chunk.detach().cpu().numpy().reshape(-1)377 378 379class StreamAudioEmitter:380 def __init__(self, sample_rate: int, prebuffer_seconds: float):381 self.sample_rate = sample_rate382 self._buffer_tracker = BufferedAudioTracker(sample_rate)383 self._prebuffer_target = max(0.0, float(prebuffer_seconds))384 self._prebuffering = self._prebuffer_target > 0.0385 self._pending_chunks: list[np.ndarray] = []386 self._pending_samples = 0387 self.chunk_count = 0388 self.has_audio = False389 390 def wait_for_capacity(self, threshold_seconds: float) -> None:391 _maybe_wait_for_buffer(self._buffer_tracker, threshold_seconds)392 393 def emit_many(self, chunks: Iterator[np.ndarray], message_prefix: str) -> Iterator[StreamEvent]:394 for chunk in chunks:395 yield from self.emit(chunk, message_prefix)396 397 def emit(self, chunk: np.ndarray, message_prefix: str) -> Iterator[StreamEvent]:398 chunk = np.asarray(chunk).reshape(-1)399 if chunk.size == 0:400 return401 if self._prebuffering:402 self._pending_chunks.append(chunk)403 self._pending_samples += int(chunk.size)404 if (self._pending_samples / self.sample_rate) < self._prebuffer_target:405 return406 self._prebuffering = False407 pending_chunks = self._pending_chunks408 self._pending_chunks = []409 self._pending_samples = 0410 for pending in pending_chunks:411 yield self._make_event(pending, message_prefix)412 return413 yield self._make_event(chunk, message_prefix)414 415 def flush(self, message_prefix: str) -> Iterator[StreamEvent]:416 if not self._prebuffering or not self._pending_chunks:417 self._prebuffering = False418 return419 self._prebuffering = False420 pending_chunks = self._pending_chunks421 self._pending_chunks = []422 self._pending_samples = 0423 for chunk in pending_chunks:424 yield self._make_event(chunk, message_prefix)425 426 def _make_event(self, chunk: np.ndarray, message_prefix: str) -> StreamEvent:427 self.chunk_count += 1428 self.has_audio = True429 self._buffer_tracker.add_chunk(chunk)430 return StreamEvent(431 message=f"{message_prefix} chunk {self.chunk_count}",432 audio=(self.sample_rate, chunk),433 )434 435 436def _maybe_wait_for_buffer(buffer_tracker: BufferedAudioTracker, threshold_seconds: float) -> None:437 if threshold_seconds <= 0:438 return439 while buffer_tracker.buffered_seconds() > threshold_seconds:440 time.sleep(0.01)441 442 443def _sanitize_tokens(444 tokens: torch.Tensor,445 codebook_size: int,446 audio_eos_token: int,447) -> tuple[torch.Tensor, bool]:448 if tokens.dim() == 1:449 tokens = tokens.unsqueeze(0)450 if tokens.numel() == 0:451 return tokens, False452 eos_rows = (tokens[:, 0] == audio_eos_token).nonzero(as_tuple=False)453 invalid_rows = ((tokens < 0) | (tokens >= codebook_size)).any(dim=1)454 stop_idx = None455 if eos_rows.numel() > 0:456 stop_idx = int(eos_rows[0].item())457 if invalid_rows.any():458 invalid_idx = int(invalid_rows.nonzero(as_tuple=False)[0].item())459 stop_idx = invalid_idx if stop_idx is None else min(stop_idx, invalid_idx)460 if stop_idx is not None:461 tokens = tokens[:stop_idx]462 return tokens, True463 return tokens, False464 465 466def _build_streaming_session(467 model: MossTTSRealtime,468 tokenizer,469 processor: MossTTSRealtimeProcessor,470 codec,471 *,472 max_length: int,473 chunk_duration: float,474 temperature: float,475 top_p: float,476 top_k: int,477 do_sample: bool,478 repetition_penalty: float,479 repetition_window: int,480) -> tuple[MossTTSRealtimeStreamingSession, MossTTSRealtimeInference]:481 inferencer = MossTTSRealtimeInference(model, tokenizer, max_length=max_length)482 inferencer.reset_generation_state(keep_cache=False)483 session = MossTTSRealtimeStreamingSession(484 inferencer,485 processor,486 codec=codec,487 codec_sample_rate=SAMPLE_RATE,488 codec_encode_kwargs={"chunk_duration": chunk_duration},489 prefill_text_len=processor.delay_tokens_len,490 temperature=temperature,491 top_p=top_p,492 top_k=top_k,493 do_sample=do_sample,494 repetition_penalty=repetition_penalty,495 repetition_window=repetition_window,496 )497 return session, inferencer498 499 500def _build_frame_decoder(501 codec,502 inferencer: MossTTSRealtimeInference,503 device: torch.device,504 *,505 chunk_frames: int,506 overlap_frames: int,507) -> AudioFrameDecoder:508 decoder = AudioStreamDecoder(509 codec,510 chunk_frames=chunk_frames,511 overlap_frames=overlap_frames,512 decode_kwargs={"chunk_duration": -1},513 device=device,514 )515 return AudioFrameDecoder(516 decoder,517 int(getattr(codec, "codebook_size", 1024)),518 int(getattr(inferencer, "audio_eos_token", 1026)),519 )520 521 522def _normalize_seed(value: float | int | None) -> int | None:523 if value is None:524 return None525 seed = int(value)526 return None if seed == 0 else seed527 528 529def _format_completion_status(530 chunk_count: int,531 sample_rate: int,532 full_audio: np.ndarray,533 started_at: float,534 first_audio_time: float | None,535) -> str:536 elapsed = time.monotonic() - started_at537 audio_seconds = float(full_audio.size) / float(sample_rate) if full_audio.size > 0 else 0.0538 rtf = (elapsed / audio_seconds) if audio_seconds > 0 else float("inf")539 parts = [540 "Done",541 ]542 return " | ".join(parts)543 544 545@functools.lru_cache(maxsize=1)546def _load_backend(547 model_path: str,548 tokenizer_path: str,549 codec_model_path: str,550 device_str: str,551 attn_impl: str,552):553 # ZeroGPU: do not call torch.cuda.is_available() here; it may trigger low-level CUDA init.554 device = torch.device(device_str)555 tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)556 processor = MossTTSRealtimeProcessor(tokenizer)557 558 # ZeroGPU: avoid torch.cuda.is_bf16_supported() before CUDA is fully managed.559 dtype = torch.bfloat16560 if attn_impl and attn_impl.lower() not in {"none", ""}:561 model = MossTTSRealtime.from_pretrained(model_path, attn_implementation=attn_impl, torch_dtype=dtype).to(device)562 if (563 attn_impl.lower() == "flash_attention_2"564 and hasattr(model, "language_model")565 and hasattr(model.language_model, "config")566 ):567 model.language_model.config.attn_implementation = "flash_attention_2"568 else:569 model = MossTTSRealtime.from_pretrained(model_path, torch_dtype=dtype).to(device)570 model.eval()571 572 codec = _load_codec(device, codec_model_path)573 return model, tokenizer, processor, codec, device574 575 576def _resolve_audio_path(audio_path: str | None, use_default: bool, default_path: str | Path) -> Path | None:577 if audio_path:578 return Path(audio_path).expanduser()579 if use_default:580 return Path(default_path).expanduser()581 return None582 583 584class StreamingTTSDemo:585 def __init__(self, audio_token_cache_size: int = 8):586 self._audio_token_cache_size = max(1, int(audio_token_cache_size))587 self._audio_token_cache: OrderedDict[tuple[str, int, float], np.ndarray] = OrderedDict()588 589 def get_or_load_backend(self, backend: BackendPaths):590 return _load_backend(591 backend.model_path,592 backend.tokenizer_path,593 backend.codec_model_path,594 backend.device_str,595 backend.attn_impl,596 )597 598 def _validate_request(self, request: StreamingRequest) -> tuple[Path | None, Path | None]:599 if not request.user_text.strip():600 raise ValueError("user_text is required.")601 if not request.assistant_text.strip():602 raise ValueError("assistant_text is required.")603 if request.streaming.text_chunk_tokens <= 0:604 raise ValueError("text_chunk_tokens must be greater than 0.")605 if request.streaming.decode_chunk_frames <= 0:606 raise ValueError("decode_chunk_frames must be greater than 0.")607 if request.streaming.chunk_duration <= 0:608 raise ValueError("chunk_duration must be greater than 0.")609 610 prompt_path = _resolve_audio_path(request.prompt_audio, request.use_default_prompt, PROMPT_WAV)611 user_path = _resolve_audio_path(request.user_audio, request.use_default_user, USER_WAV)612 613 if prompt_path is not None and not prompt_path.exists():614 raise FileNotFoundError(f"Prompt wav not found: {prompt_path}")615 if user_path is not None and not user_path.exists():616 raise FileNotFoundError(f"User wav not found: {user_path}")617 618 return prompt_path, user_path619 620 def _encode_audio_tokens(621 self,622 path: Path,623 codec,624 device: torch.device,625 chunk_duration: float,626 ) -> np.ndarray:627 resolved_path = path.expanduser().resolve()628 cache_key = (str(resolved_path), int(resolved_path.stat().st_mtime_ns), float(chunk_duration))629 cached_tokens = self._audio_token_cache.get(cache_key)630 if cached_tokens is not None:631 self._audio_token_cache.move_to_end(cache_key)632 return cached_tokens633 634 with torch.inference_mode():635 audio_tensor = _load_audio(resolved_path)636 waveform = audio_tensor.to(device)637 if waveform.dim() == 2:638 waveform = waveform.unsqueeze(0)639 encode_result = codec.encode(waveform, chunk_duration=chunk_duration)640 641 tokens = _extract_codes(encode_result)642 if isinstance(tokens, torch.Tensor):643 tokens = tokens.detach().cpu().numpy()644 else:645 tokens = np.asarray(tokens)646 647 self._audio_token_cache[cache_key] = tokens648 self._audio_token_cache.move_to_end(cache_key)649 while len(self._audio_token_cache) > self._audio_token_cache_size:650 self._audio_token_cache.popitem(last=False)651 652 return tokens653 654 @staticmethod655 def _build_text_only_turn_input(656 processor: MossTTSRealtimeProcessor,657 user_text: str,658 prompt_tokens: np.ndarray | None,659 ) -> np.ndarray:660 system_prompt = processor.make_ensemble(prompt_tokens)661 user_prompt_text = "<|im_end|>\n<|im_start|>user\n" + user_text + "<|im_end|>\n<|im_start|>assistant\n"662 user_prompt_tokens = processor.tokenizer(user_prompt_text)["input_ids"]663 user_prompt = np.full(664 shape=(len(user_prompt_tokens), processor.channels + 1),665 fill_value=processor.audio_channel_pad,666 dtype=np.int64,667 )668 user_prompt[:, 0] = np.asarray(user_prompt_tokens, dtype=np.int64)669 return np.concatenate([system_prompt, user_prompt], axis=0)670 671 def _prepare_session_turn(672 self,673 session: MossTTSRealtimeStreamingSession,674 processor: MossTTSRealtimeProcessor,675 user_text: str,676 prompt_tokens: np.ndarray | None,677 user_tokens: np.ndarray | None,678 ) -> str | None:679 if user_tokens is None:680 turn_input_ids = self._build_text_only_turn_input(processor, user_text, prompt_tokens)681 session.reset_turn(input_ids=turn_input_ids, include_system_prompt=True, reset_cache=True)682 return "No user audio provided, running text-only turn."683 684 session.reset_turn(685 user_text=user_text,686 user_audio_tokens=user_tokens,687 include_system_prompt=True,688 reset_cache=True,689 )690 return None691 692 def run_stream(self, request: StreamingRequest) -> Iterator[StreamEvent]:693 prompt_path, user_path = self._validate_request(request)694 model, tokenizer, processor, codec, device = self.get_or_load_backend(request.backend)695 _apply_seed(request.generation.seed)696 697 prompt_tokens = (698 self._encode_audio_tokens(699 prompt_path,700 codec,701 device,702 chunk_duration=request.streaming.chunk_duration,703 )704 if prompt_path is not None705 else None706 )707 user_tokens = (708 self._encode_audio_tokens(709 user_path,710 codec,711 device,712 chunk_duration=request.streaming.chunk_duration,713 )714 if user_path is not None715 else None716 )717 718 session, inferencer = _build_streaming_session(719 model,720 tokenizer,721 processor,722 codec,723 max_length=request.generation.max_length,724 chunk_duration=request.streaming.chunk_duration,725 temperature=request.generation.temperature,726 top_p=request.generation.top_p,727 top_k=request.generation.top_k,728 do_sample=request.generation.do_sample,729 repetition_penalty=request.generation.repetition_penalty,730 repetition_window=request.generation.repetition_window,731 )732 if prompt_tokens is not None:733 session.set_voice_prompt_tokens(prompt_tokens)734 else:735 session.clear_voice_prompt()736 737 turn_message = self._prepare_session_turn(738 session,739 processor,740 request.user_text,741 prompt_tokens,742 user_tokens,743 )744 if turn_message:745 yield StreamEvent(message=turn_message)746 747 frame_decoder = _build_frame_decoder(748 codec,749 inferencer,750 device,751 chunk_frames=request.streaming.decode_chunk_frames,752 overlap_frames=request.streaming.decode_overlap_frames,753 )754 755 text_tokens = tokenizer.encode(request.assistant_text, add_special_tokens=False)756 if not text_tokens:757 raise RuntimeError("Assistant text tokenization returned no tokens.")758 759 token_stream = TokenChunkStream(text_tokens, request.streaming.text_chunk_tokens)760 audio_emitter = StreamAudioEmitter(SAMPLE_RATE, request.streaming.prebuffer_seconds)761 762 with codec.streaming(batch_size=1):763 for token_chunk in token_stream:764 audio_emitter.wait_for_capacity(request.streaming.buffer_threshold_seconds)765 audio_frames = session.push_text_tokens(token_chunk)766 yield from audio_emitter.emit_many(frame_decoder.decode_frames(audio_frames), "Streaming")767 if request.streaming.input_delay > 0:768 time.sleep(request.streaming.input_delay)769 770 final_frames = session.end_text()771 yield from audio_emitter.emit_many(frame_decoder.decode_frames(final_frames), "Finalizing")772 773 while True:774 drain_frames = session.drain(max_steps=1)775 if not drain_frames:776 break777 yield from audio_emitter.emit_many(frame_decoder.decode_frames(drain_frames), "Finalizing")778 if session.inferencer.is_finished:779 break780 781 yield from audio_emitter.emit_many(frame_decoder.flush(), "Final")782 yield from audio_emitter.flush("Final")783 784 if not audio_emitter.has_audio:785 raise RuntimeError("No audio waveform chunks decoded from streaming inference.")786 787 yield StreamEvent(message="Streaming complete.")788 789 790class WarmupManager:791 def __init__(self, tts_demo: "StreamingTTSDemo", backend: BackendPaths):792 self.tts_demo = tts_demo793 self.backend = backend794 self._lock = threading.Lock()795 self._thread: threading.Thread | None = None796 self._started = False797 # ZeroGPU: startup warmup is disabled because it initializes CUDA outside @spaces.GPU.798 self._state = "ready"799 self._progress = 1.0800 self._message = "Ready."801 self._detail = "Startup warmup disabled for ZeroGPU; the first generation will load the model."802 self._error: str | None = None803 804 def start(self) -> None:805 with self._lock:806 if self._started:807 return808 self._started = True809 self._thread = threading.Thread(target=self._run, name="tts-startup-warmup", daemon=True)810 self._thread.start()811 812 def snapshot(self) -> WarmupSnapshot:813 with self._lock:814 return WarmupSnapshot(815 state=self._state,816 progress=self._progress,817 message=self._message,818 detail=self._detail,819 error=self._error,820 )821 822 def _set_state(823 self,824 *,825 state: str | None = None,826 progress: float | None = None,827 message: str | None = None,828 detail: str | None = None,829 error: str | None = None,830 ) -> None:831 with self._lock:832 if state is not None:833 self._state = state834 if progress is not None:835 self._progress = max(0.0, min(1.0, float(progress)))836 if message is not None:837 self._message = message838 if detail is not None:839 self._detail = detail840 self._error = error841 842 @staticmethod843 def _consume_audio(chunks: Iterator[np.ndarray]) -> None:844 for _chunk in chunks:845 pass846 847 @staticmethod848 def _ensure_warmup_text(tokenizer, minimum_tokens: int) -> tuple[str, list[int]]:849 text = WARMUP_BASE_ASSISTANT_TEXT850 tokens = tokenizer.encode(text, add_special_tokens=False)851 while len(tokens) < minimum_tokens:852 text = f"{text} {WARMUP_BASE_ASSISTANT_TEXT}"853 tokens = tokenizer.encode(text, add_special_tokens=False)854 return text, tokens855 856 @staticmethod857 def _warmup_step_detail(step_idx: int, total_steps: int) -> str:858 if step_idx == 1:859 return "First incremental step is compiling the cold streaming path."860 if step_idx == 2:861 return "Second incremental step is warming the next steady-state path."862 if step_idx == DEFAULT_REPETITION_WINDOW:863 return "Warming the first full repetition-window step."864 if step_idx == WARMUP_STEP_TOKENS:865 return "Confirming the post-window steady-state step."866 return f"Warming token step {step_idx}/{total_steps}."867 868 def _run(self) -> None:869 try:870 self._set_state(871 state="running",872 progress=0.02,873 message="Starting startup warmup.",874 detail="Preparing backend state for the first real request.",875 error=None,876 )877 878 self._set_state(879 progress=0.08,880 message="Loading backend.",881 detail="Model, tokenizer, codec, and CUDA runtime are warming up.",882 error=None,883 )884 model, tokenizer, processor, codec, device = self.tts_demo.get_or_load_backend(self.backend)885 886 self._set_state(887 progress=0.32,888 message="Preparing streaming session.",889 detail="Building a text-only warmup turn and its decoder.",890 error=None,891 )892 session, inferencer = _build_streaming_session(893 model,894 tokenizer,895 processor,896 codec,897 max_length=256,898 chunk_duration=0.24,899 temperature=0.8,900 top_p=0.6,901 top_k=30,902 do_sample=True,903 repetition_penalty=1.1,904 repetition_window=DEFAULT_REPETITION_WINDOW,905 )906 session.clear_voice_prompt()907 session.reset_turn(908 input_ids=self.tts_demo._build_text_only_turn_input(processor, WARMUP_USER_TEXT, None),909 include_system_prompt=True,910 reset_cache=True,911 )912 913 frame_decoder = _build_frame_decoder(914 codec,915 inferencer,916 device,917 chunk_frames=WARMUP_STEP_TOKENS,918 overlap_frames=0,919 )920 921 _, warmup_tokens = self._ensure_warmup_text(922 tokenizer,923 processor.delay_tokens_len + WARMUP_STEP_TOKENS,924 )925 926 with codec.streaming(batch_size=1):927 self._set_state(928 progress=0.45,929 message="Running prefill.",930 detail="Building the first KV cache and warming the backbone path.",931 error=None,932 )933 prefill_frames = session.push_text_tokens(warmup_tokens[: processor.delay_tokens_len])934 self._consume_audio(frame_decoder.decode_frames(prefill_frames))935 936 step_tokens = warmup_tokens[937 processor.delay_tokens_len : processor.delay_tokens_len + WARMUP_STEP_TOKENS938 ]939 total_steps = max(1, len(step_tokens))940 for idx, token in enumerate(step_tokens, start=1):941 self._set_state(942 progress=0.55 + 0.25 * (idx - 1) / total_steps,943 message="Compiling first streaming steps.",944 detail=self._warmup_step_detail(idx, total_steps),945 error=None,946 )947 step_frames = session.push_text_tokens([token])948 self._consume_audio(frame_decoder.decode_frames(step_frames))949 950 self._set_state(951 progress=0.86,952 message="Warming finalization path.",953 detail="Priming end-text, drain, and decoder flush before user traffic.",954 error=None,955 )956 final_frames = session.end_text()957 self._consume_audio(frame_decoder.decode_frames(final_frames))958 drain_frames = session.drain(max_steps=1)959 self._consume_audio(frame_decoder.decode_frames(drain_frames))960 self._consume_audio(frame_decoder.flush())961 962 self._set_state(963 state="ready",964 progress=1.0,965 message="Warmup complete.",966 detail="The first real request should avoid the cold-start stall.",967 error=None,968 )969 except Exception as exc:970 self._set_state(971 state="failed",972 progress=1.0,973 message="Warmup failed.",974 detail="The app did not finish startup warmup.",975 error=str(exc),976 )977 print(f"[MossTTSRealtime][warmup-error] {exc}", file=sys.stderr, flush=True)978 979 980def _warmup_button_update(snapshot: WarmupSnapshot):981 if snapshot.ready:982 return gr.update(value="Generate", interactive=True)983 if snapshot.failed:984 return gr.update(value="Warmup Failed", interactive=False)985 return gr.update(value="Warming Up...", interactive=False)986 987 988def _warmup_gate_message(snapshot: WarmupSnapshot) -> str:989 progress_pct = int(round(max(0.0, min(1.0, snapshot.progress)) * 100.0))990 if snapshot.failed:991 return f"Warmup failed: {snapshot.error or snapshot.message}"992 return f"Warmup in progress ({progress_pct}%): {snapshot.message}"993 994 995def _status_from_snapshot(snapshot: WarmupSnapshot) -> str:996 return "Ready." if snapshot.ready else _warmup_gate_message(snapshot)997 998 999def _warmup_status_update(snapshot: WarmupSnapshot):1000 return gr.update(value=_status_from_snapshot(snapshot))1001 1002 1003def _warmup_timer_update(snapshot: WarmupSnapshot):1004 return gr.update(active=not (snapshot.ready or snapshot.failed))1005 1006 1007def _encode_chunk(sr: int, chunk: np.ndarray, idx: int) -> str:1008 if chunk.dtype != np.float32:1009 chunk = chunk.astype(np.float32)1010 if chunk.ndim != 1:1011 chunk = chunk.reshape(-1)1012 payload = {1013 "sr": int(sr),1014 "idx": int(idx),1015 "data": base64.b64encode(chunk.tobytes()).decode("ascii"),1016 }1017 return json.dumps(payload)1018 1019 1020def _build_request(1021 args: argparse.Namespace,1022 *,1023 user_text: str | None,1024 assistant_text: str | None,1025 prompt_audio: str | None,1026 user_audio: str | None,1027 use_default_prompt: bool,1028 use_default_user: bool,1029 temperature: float,1030 top_p: float,1031 top_k: int,1032 repetition_penalty: float,1033 repetition_window: int,1034 do_sample: bool,1035 max_length: int,1036 seed: float | int | None,1037 text_chunk_tokens: int,1038 input_delay: float,1039 decode_chunk_frames: int,1040 decode_overlap_frames: int,1041 chunk_duration: float,1042 prebuffer_seconds: float,1043) -> StreamingRequest:1044 return StreamingRequest(1045 user_text=str(user_text or "Hello!"),1046 assistant_text=str(assistant_text or ""),1047 prompt_audio=prompt_audio,1048 user_audio=user_audio,1049 use_default_prompt=use_default_prompt,1050 use_default_user=use_default_user,1051 generation=GenerationConfig(1052 temperature=float(temperature),1053 top_p=float(top_p),1054 top_k=int(top_k),1055 repetition_penalty=float(repetition_penalty),1056 repetition_window=int(repetition_window),1057 do_sample=bool(do_sample),1058 max_length=int(max_length),1059 seed=_normalize_seed(seed),1060 ),1061 streaming=StreamingConfig(1062 text_chunk_tokens=int(text_chunk_tokens),1063 input_delay=float(input_delay),1064 decode_chunk_frames=int(decode_chunk_frames),1065 decode_overlap_frames=int(decode_overlap_frames),1066 chunk_duration=float(chunk_duration),1067 prebuffer_seconds=float(prebuffer_seconds),1068 ),1069 backend=BackendPaths(1070 model_path=args.model_path,1071 tokenizer_path=args.tokenizer_path,1072 codec_model_path=args.codec_model_path,1073 device_str=args.device,1074 attn_impl=args.attn_implementation,1075 ),1076 )1077 1078 1079STREAM_PLAYER_HTML = """1080<style>1081#pcm_stream {1082 position: absolute !important;1083 left: -9999px !important;1084 width: 1px !important;1085 height: 1px !important;1086 opacity: 0 !important;1087 pointer-events: none !important;1088}1089#pcm_stream textarea, #pcm_stream input {1090 width: 1px !important;1091 height: 1px !important;1092 opacity: 0 !important;1093}1094</style>1095"""1096 1097STREAM_PLAYER_JS = r"""1098const elemId = "pcm_stream";1099if (window.__pcm_streaming_inited__) {1100 return;1101}1102window.__pcm_streaming_inited__ = true;1103 1104let audioCtx = null;1105let nextTime = 0;1106let lastIdx = -1;1107let lastValue = "";1108let boundField = null;1109let usingSetterHook = false;1110const FADE_MS = 6;1111const MIN_BUFFER_SEC = 0.25;1112 1113function initAudio(sr) {1114 if (audioCtx && audioCtx.sampleRate !== sr) {1115 audioCtx.close();1116 audioCtx = null;1117 }1118 if (!audioCtx) {1119 audioCtx = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: sr });1120 nextTime = audioCtx.currentTime;1121 }1122 if (audioCtx.state === "suspended") {1123 audioCtx.resume();1124 }1125}1126 1127function decodeBase64ToFloat32(base64) {1128 const binary = atob(base64);1129 const len = binary.length;1130 const bytes = new Uint8Array(len);1131 for (let i = 0; i < len; i++) {1132 bytes[i] = binary.charCodeAt(i);1133 }1134 return new Float32Array(bytes.buffer);1135}1136 1137function playChunk(samples, sr, idx) {1138 initAudio(sr);1139 const buffer = audioCtx.createBuffer(1, samples.length, sr);1140 buffer.copyToChannel(samples, 0);1141 const source = audioCtx.createBufferSource();1142 source.buffer = buffer;1143 const gain = audioCtx.createGain();1144 source.connect(gain);1145 gain.connect(audioCtx.destination);1146 const now = audioCtx.currentTime;1147 if (nextTime < now + MIN_BUFFER_SEC) {1148 nextTime = now + MIN_BUFFER_SEC;1149 }1150 const startTime = Math.max(now, nextTime);1151 const endTime = startTime + buffer.duration;1152 const fade = Math.min(FADE_MS / 1000.0, buffer.duration / 4);1153 gain.gain.setValueAtTime(0.0, startTime);1154 gain.gain.linearRampToValueAtTime(1.0, startTime + fade);1155 gain.gain.setValueAtTime(1.0, Math.max(startTime + fade, endTime - fade));1156 gain.gain.linearRampToValueAtTime(0.0, endTime);1157 source.start(startTime);1158 nextTime = endTime;1159}1160 1161function handlePayload(text) {1162 if (!text) return;1163 let payload;1164 try {1165 payload = JSON.parse(text);1166 } catch (e) {1167 return;1168 }1169 if (Array.isArray(payload)) {1170 for (const item of payload) {1171 handlePayloadObject(item);1172 }1173 return;1174 }1175 handlePayloadObject(payload);1176}1177 1178function handlePayloadObject(payload) {1179 if (!payload) return;1180 if (payload.reset) {1181 lastIdx = -1;1182 lastValue = "";1183 if (audioCtx) {1184 audioCtx.close();1185 audioCtx = null;1186 }1187 return;1188 }1189 const idx = payload.idx ?? 0;1190 if (idx <= lastIdx) return;1191 lastIdx = idx;1192 const sr = payload.sr || 24000;1193 const samples = decodeBase64ToFloat32(payload.data);1194 playChunk(samples, sr, idx);1195}1196 1197function hookField(field) {1198 if (!field || field === boundField) return;1199 boundField = field;1200 const proto = field.tagName === "TEXTAREA" ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;