mkultr4's picture
Update app.py
150feea verified
Raw
History Blame Contribute Delete
25.4 kB
"""
MOSS-TTS-Realtime — Hugging Face Space
Adapted from the official app.py with proper streaming via
MossTTSRealtimeStreamingSession + AudioStreamDecoder + Web Audio API.
"""
import base64
import functools
import importlib.util
import json
import os
import subprocess
import sys
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterator, Sequence
# ── spaces (ZeroGPU — must come before other imports) ─────────────────────────
try:
import spaces
except ImportError:
class _SpacesFallback:
@staticmethod
def GPU(*_args, **_kwargs):
def _decorator(func):
return func
return _decorator
spaces = _SpacesFallback()
# ── Repo setup ─────────────────────────────────────────────────────────────────
REPO_DIR = Path("/home/user/moss-tts")
REALTIME_DIR = REPO_DIR / "moss_tts_realtime"
def _run(*cmd, **kw):
subprocess.check_call(list(cmd), **kw)
if not REPO_DIR.exists():
print("[INFO] Cloning MOSS-TTS repository…")
_run("git", "clone", "--depth=1",
"https://github.com/OpenMOSS/MOSS-TTS.git", str(REPO_DIR))
print("[INFO] Clone complete.")
for p in [str(REALTIME_DIR), str(REPO_DIR)]:
if p not in sys.path:
sys.path.insert(0, p)
# ── Core imports ───────────────────────────────────────────────────────────────
import gradio as gr
import numpy as np
import torch
import torch._dynamo
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors = True
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") # disable compile globally — avoids torch.device ConstantVariable bug
import torchaudio
from transformers import AutoTokenizer, AutoModel
from mossttsrealtime import MossTTSRealtime, MossTTSRealtimeProcessor
from mossttsrealtime.streaming_mossttsrealtime import (
AudioStreamDecoder,
MossTTSRealtimeInference,
MossTTSRealtimeStreamingSession,
)
torch.backends.cuda.enable_cudnn_sdp(False)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
# ── Constants ──────────────────────────────────────────────────────────────────
MODEL_PATH = "OpenMOSS-Team/MOSS-TTS-Realtime"
TOKENIZER_PATH = "OpenMOSS-Team/MOSS-TTS-Realtime"
CODEC_MODEL_PATH = "OpenMOSS-Team/MOSS-Audio-Tokenizer"
SAMPLE_RATE = 24_000
SUPPORTED_LANGUAGES = (
"Chinese, English, Turkish, Korean, Japanese, German, French, "
"Spanish, Portuguese, Italian, Russian, Arabic, Hindi, "
"Indonesian, Vietnamese, Dutch, Polish, Swedish, Danish, Norwegian"
)
# ── Dataclasses (mirrors official app.py) ─────────────────────────────────────
@dataclass(frozen=True)
class GenerationConfig:
temperature: float
top_p: float
top_k: int
repetition_penalty: float
repetition_window: int
do_sample: bool
max_length: int
@dataclass(frozen=True)
class StreamingConfig:
text_chunk_tokens: int
decode_chunk_frames: int
decode_overlap_frames: int
chunk_duration: float
@dataclass(frozen=True)
class StreamEvent:
message: str
audio: "tuple[int, np.ndarray] | None" = None
# ── Backend loader ─────────────────────────────────────────────────────────────
@functools.lru_cache(maxsize=1)
def _load_backend(model_path: str, tokenizer_path: str, codec_path: str,
device_str: str, attn_impl: str):
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required.")
device = torch.device(device_str)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
processor = MossTTSRealtimeProcessor(tokenizer)
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
model = MossTTSRealtime.from_pretrained(
model_path, attn_implementation=attn_impl, torch_dtype=dtype
).to(device).eval()
# Disable torch.compile on the local transformer — torch.device causes
# a ConstantVariable assertion error in torch 2.9.1+cu128 dynamo.
codec = AutoModel.from_pretrained(
codec_path, trust_remote_code=True
).eval().to(device)
print("[INFO] Backend loaded.")
return model, tokenizer, processor, codec, device
def resolve_attn(device_str: str) -> str:
if not torch.cuda.is_available():
return "eager"
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
if (importlib.util.find_spec("flash_attn") is not None
and dtype in {torch.float16, torch.bfloat16}):
major, _ = torch.cuda.get_device_capability()
if major >= 8:
return "flash_attention_2"
return "sdpa"
# ── Audio helpers ──────────────────────────────────────────────────────────────
_audio_token_cache: OrderedDict = OrderedDict()
_CACHE_SIZE = 8
def _load_audio(path: Path) -> torch.Tensor:
wav, sr = torchaudio.load(path)
if sr != SAMPLE_RATE:
wav = torchaudio.functional.resample(wav, sr, SAMPLE_RATE)
if wav.shape[0] > 1:
wav = wav.mean(dim=0, keepdim=True)
return wav
def _extract_codes(encode_result):
if isinstance(encode_result, dict):
codes = encode_result["audio_codes"]
elif isinstance(encode_result, (list, tuple)) and encode_result:
codes = encode_result[0]
else:
codes = encode_result
if isinstance(codes, np.ndarray):
codes = torch.from_numpy(codes)
if isinstance(codes, torch.Tensor) and codes.dim() == 3:
codes = codes[:, 0, :] if codes.shape[1] == 1 else codes[0]
return codes
def _encode_audio_tokens(path: Path, codec, device: torch.device,
chunk_duration: float) -> np.ndarray:
key = (str(path.resolve()), int(path.stat().st_mtime_ns), chunk_duration)
if key in _audio_token_cache:
_audio_token_cache.move_to_end(key)
return _audio_token_cache[key]
with torch.inference_mode():
wav = _load_audio(path).to(device)
if wav.dim() == 2:
wav = wav.unsqueeze(0)
result = codec.encode(wav, chunk_duration=chunk_duration)
tokens = _extract_codes(result)
if isinstance(tokens, torch.Tensor):
tokens = tokens.detach().cpu().numpy()
_audio_token_cache[key] = tokens
_audio_token_cache.move_to_end(key)
while len(_audio_token_cache) > _CACHE_SIZE:
_audio_token_cache.popitem(last=False)
return tokens
def _sanitize_tokens(tokens: torch.Tensor, codebook_size: int,
audio_eos_token: int) -> tuple:
if tokens.dim() == 1:
tokens = tokens.unsqueeze(0)
if tokens.numel() == 0:
return tokens, False
eos_rows = (tokens[:, 0] == audio_eos_token).nonzero(as_tuple=False)
invalid_rows = ((tokens < 0) | (tokens >= codebook_size)).any(dim=1)
stop_idx = None
if eos_rows.numel() > 0:
stop_idx = int(eos_rows[0].item())
if invalid_rows.any():
inv = int(invalid_rows.nonzero(as_tuple=False)[0].item())
stop_idx = inv if stop_idx is None else min(stop_idx, inv)
if stop_idx is not None:
return tokens[:stop_idx], True
return tokens, False
def _decode_frames(audio_frames, decoder: AudioStreamDecoder,
codebook_size: int, audio_eos_token: int) -> Iterator[np.ndarray]:
for frame in audio_frames:
tokens = frame
if tokens.dim() == 3:
tokens = tokens[0]
if tokens.dim() != 2:
continue
tokens, _ = _sanitize_tokens(tokens, codebook_size, audio_eos_token)
if tokens.numel() == 0:
continue
decoder.push_tokens(tokens.detach())
for wav in decoder.audio_chunks():
if wav.numel() > 0:
yield wav.detach().cpu().numpy().reshape(-1)
def _flush_decoder(decoder: AudioStreamDecoder) -> Iterator[np.ndarray]:
final = decoder.flush()
if final is not None and final.numel() > 0:
yield final.detach().cpu().numpy().reshape(-1)
def _encode_chunk(sr: int, chunk: np.ndarray, idx: int) -> str:
chunk = chunk.astype(np.float32).reshape(-1)
return json.dumps({
"sr": int(sr), "idx": int(idx),
"data": base64.b64encode(chunk.tobytes()).decode("ascii"),
})
# ── Core streaming pipeline ────────────────────────────────────────────────────
def _build_text_only_turn(processor, user_text: str,
prompt_tokens) -> np.ndarray:
system_prompt = processor.make_ensemble(prompt_tokens)
user_text_str = "<|im_end|>\n<|im_start|>user\n" + user_text + \
"<|im_end|>\n<|im_start|>assistant\n"
user_tok = processor.tokenizer(user_text_str)["input_ids"]
user_arr = np.full((len(user_tok), processor.channels + 1),
fill_value=processor.audio_channel_pad, dtype=np.int64)
user_arr[:, 0] = np.asarray(user_tok, dtype=np.int64)
return np.concatenate([system_prompt, user_arr], axis=0)
def run_stream(
assistant_text: str,
prompt_audio_path: "str | None",
user_text: str,
gen: GenerationConfig,
streaming: StreamingConfig,
attn_impl: str,
) -> Iterator[StreamEvent]:
model, tokenizer, processor, codec, device = _load_backend(
MODEL_PATH, TOKENIZER_PATH, CODEC_MODEL_PATH, "cuda:0", attn_impl
)
prompt_tokens = None
if prompt_audio_path and Path(prompt_audio_path).exists():
prompt_tokens = _encode_audio_tokens(
Path(prompt_audio_path), codec, device, streaming.chunk_duration
)
inferencer = MossTTSRealtimeInference(
model, tokenizer, max_length=gen.max_length
)
inferencer.reset_generation_state(keep_cache=False)
session = MossTTSRealtimeStreamingSession(
inferencer, processor,
codec=codec,
codec_sample_rate=SAMPLE_RATE,
codec_encode_kwargs={"chunk_duration": streaming.chunk_duration},
prefill_text_len=processor.delay_tokens_len,
temperature=gen.temperature,
top_p=gen.top_p,
top_k=gen.top_k,
do_sample=gen.do_sample,
repetition_penalty=gen.repetition_penalty,
repetition_window=gen.repetition_window,
)
if prompt_tokens is not None:
session.set_voice_prompt_tokens(prompt_tokens)
else:
session.clear_voice_prompt()
turn_input = _build_text_only_turn(processor, user_text, prompt_tokens)
session.reset_turn(
input_ids=turn_input, include_system_prompt=True, reset_cache=True
)
decoder = AudioStreamDecoder(
codec,
chunk_frames=streaming.decode_chunk_frames,
overlap_frames=streaming.decode_overlap_frames,
decode_kwargs={"chunk_duration": -1},
device=device,
)
codebook_size = int(getattr(codec, "codebook_size", 1024))
audio_eos_token = int(getattr(inferencer, "audio_eos_token", 1026))
text_tokens = tokenizer.encode(assistant_text, add_special_tokens=False)
if not text_tokens:
raise RuntimeError("No tokens from assistant text.")
chunk_size = max(1, streaming.text_chunk_tokens)
token_chunks = [text_tokens[i:i+chunk_size]
for i in range(0, len(text_tokens), chunk_size)]
with codec.streaming(batch_size=1):
for chunk in token_chunks:
audio_frames = session.push_text_tokens(chunk)
yield from (
StreamEvent(message="Streaming", audio=(SAMPLE_RATE, w))
for w in _decode_frames(audio_frames, decoder,
codebook_size, audio_eos_token)
)
audio_frames = session.end_text()
yield from (
StreamEvent(message="Finalizing", audio=(SAMPLE_RATE, w))
for w in _decode_frames(audio_frames, decoder,
codebook_size, audio_eos_token)
)
while True:
audio_frames = session.drain(max_steps=1)
if not audio_frames:
break
yield from (
StreamEvent(message="Draining", audio=(SAMPLE_RATE, w))
for w in _decode_frames(audio_frames, decoder,
codebook_size, audio_eos_token)
)
if session.inferencer.is_finished:
break
yield from (
StreamEvent(message="Flushing", audio=(SAMPLE_RATE, w))
for w in _flush_decoder(decoder)
)
yield StreamEvent(message="Done")
# ── Gradio GPU wrapper ─────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def generate(
user_text: str,
assistant_text: str,
prompt_audio,
temperature: float, top_p: float, top_k: int,
repetition_penalty: float, repetition_window: int,
text_chunk_tokens: int,
decode_chunk_frames: int, decode_overlap_frames: int,
chunk_duration: float,
):
if not assistant_text.strip():
gr.Warning("Please enter assistant text to synthesize.")
return
attn_impl = resolve_attn("cuda:0")
gen = GenerationConfig(
temperature=temperature, top_p=top_p, top_k=int(top_k),
repetition_penalty=repetition_penalty,
repetition_window=int(repetition_window),
do_sample=True, max_length=3000,
)
streaming = StreamingConfig(
text_chunk_tokens=int(text_chunk_tokens),
decode_chunk_frames=int(decode_chunk_frames),
decode_overlap_frames=int(decode_overlap_frames),
chunk_duration=float(chunk_duration),
)
started_at = time.monotonic()
first_chunk_at = None
full_chunks = []
chunk_index = 0
stream_reset = json.dumps({"reset": True})
yield stream_reset, None, "Starting…"
try:
for event in run_stream(
assistant_text=assistant_text,
prompt_audio_path=prompt_audio,
user_text=user_text or "Hello!",
gen=gen, streaming=streaming, attn_impl=attn_impl,
):
if event.audio is None:
yield gr.update(), gr.update(), event.message
continue
sr, chunk = event.audio
chunk = np.asarray(chunk, dtype=np.float32).reshape(-1)
if chunk.size == 0:
continue
if first_chunk_at is None:
first_chunk_at = time.monotonic()
full_chunks.append(chunk)
chunk_index += 1
ttfb_ms = (first_chunk_at - started_at) * 1000 if first_chunk_at else float("nan")
payload = _encode_chunk(sr, chunk, chunk_index)
yield payload, gr.update(), f"Streaming… chunks={chunk_index} | TTFB={ttfb_ms:.0f}ms"
if full_chunks:
full_audio = np.concatenate(full_chunks)
elapsed = time.monotonic() - started_at
audio_sec = full_audio.size / SAMPLE_RATE
rtf = elapsed / audio_sec if audio_sec > 0 else float("inf")
ttfb_ms = (first_chunk_at - started_at) * 1000 if first_chunk_at else float("nan")
done_msg = (
f"Done | chunks={chunk_index} | audio={audio_sec:.2f}s | "
f"elapsed={elapsed:.2f}s | RTF={rtf:.3f} | TTFB={ttfb_ms:.0f}ms"
f"{' ✅' if rtf < 1.0 else ' ⚠️'}"
)
yield gr.update(), (SAMPLE_RATE, full_audio), done_msg
else:
yield gr.update(), gr.update(), "Done | no audio chunks emitted"
except Exception as exc:
import traceback; traceback.print_exc()
yield gr.update(), gr.update(), f"Error: {exc}"
# ── Gradio UI ──────────────────────────────────────────────────────────────────
STREAM_PLAYER_HTML = """
<style>
#pcm_stream {
position: absolute !important; left: -9999px !important;
width: 1px !important; height: 1px !important;
opacity: 0 !important; pointer-events: none !important;
}
#pcm_stream textarea, #pcm_stream input {
width: 1px !important; height: 1px !important; opacity: 0 !important;
}
</style>
<div id="pcm-stream-status" style="font-size:12px;color:#555;">
Live playback via Web Audio API. Click Generate to unlock audio.
</div>
<div id="pcm-stream-meta" style="font-size:12px;color:#333;margin:6px 0;">
Now Playing Chunk: <span id="pcm-stream-playing">-</span> &nbsp;|&nbsp;
Last Yielded: <span id="pcm-stream-yielded">-</span>
</div>
"""
STREAM_PLAYER_JS = r"""
const elemId = "pcm_stream";
if (window.__pcm_streaming_inited__) return;
window.__pcm_streaming_inited__ = true;
let audioCtx = null, nextTime = 0, lastIdx = -1, lastValue = "",
boundField = null, usingSetterHook = false;
const FADE_MS = 6, MIN_BUFFER_SEC = 0.25;
const statusEl = document.getElementById("pcm-stream-status");
const playingEl = document.getElementById("pcm-stream-playing");
const yieldedEl = document.getElementById("pcm-stream-yielded");
function setStatus(m){ if(statusEl) statusEl.textContent = m; }
function setPlaying(i){ if(playingEl) playingEl.textContent = `${i}`; }
function setYielded(i){ if(yieldedEl) yieldedEl.textContent = `${i}`; }
function initAudio(sr){
if(audioCtx && audioCtx.sampleRate !== sr){ audioCtx.close(); audioCtx = null; }
if(!audioCtx){ audioCtx = new (window.AudioContext||window.webkitAudioContext)({sampleRate:sr}); nextTime = audioCtx.currentTime; }
if(audioCtx.state === "suspended") audioCtx.resume();
}
function decodeBase64ToFloat32(b64){
const bin = atob(b64), len = bin.length, bytes = new Uint8Array(len);
for(let i=0;i<len;i++) bytes[i] = bin.charCodeAt(i);
return new Float32Array(bytes.buffer);
}
function playChunk(samples, sr, idx){
initAudio(sr);
const buf = audioCtx.createBuffer(1, samples.length, sr);
buf.copyToChannel(samples, 0);
const src = audioCtx.createBufferSource(), gain = audioCtx.createGain();
src.buffer = buf; src.connect(gain); gain.connect(audioCtx.destination);
const now = audioCtx.currentTime;
if(nextTime < now + MIN_BUFFER_SEC) nextTime = now + MIN_BUFFER_SEC;
const t0 = Math.max(now, nextTime), t1 = t0 + buf.duration;
const fade = Math.min(FADE_MS/1000, buf.duration/4);
gain.gain.setValueAtTime(0, t0);
gain.gain.linearRampToValueAtTime(1, t0+fade);
gain.gain.setValueAtTime(1, Math.max(t0+fade, t1-fade));
gain.gain.linearRampToValueAtTime(0, t1);
src.start(t0); nextTime = t1;
setPlaying(idx); setStatus(`Streaming… (chunk ${idx})`);
}
function handlePayloadObject(p){
if(!p) return;
if(p.reset){
lastIdx=-1; lastValue="";
if(audioCtx){ audioCtx.close(); audioCtx=null; }
setPlaying("-"); setYielded("-");
setStatus("Live playback via Web Audio API. Click Generate to unlock audio.");
return;
}
const idx = p.idx ?? 0;
if(idx <= lastIdx) return;
lastIdx = idx;
setYielded(idx);
playChunk(decodeBase64ToFloat32(p.data), p.sr||24000, idx);
}
function handlePayload(text){
if(!text) return;
let p; try{ p = JSON.parse(text); } catch(e){ return; }
if(Array.isArray(p)){ p.forEach(handlePayloadObject); } else { handlePayloadObject(p); }
}
function hookField(field){
if(!field || field === boundField) return;
boundField = field;
const proto = field.tagName==="TEXTAREA" ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype;
const desc = Object.getOwnPropertyDescriptor(proto, "value");
if(!desc||!desc.get||!desc.set){ usingSetterHook=false; return; }
usingSetterHook = true;
const nGet=desc.get, nSet=desc.set;
Object.defineProperty(field,"value",{
configurable:true,
get(){ return nGet.call(field); },
set(v){ nSet.call(field,v); if(v&&v!==lastValue){ lastValue=v; handlePayload(v); } }
});
const init = field.value;
if(init&&init!==lastValue){ lastValue=init; handlePayload(init); }
}
function pollField(){
const f = document.querySelector(`#${elemId} textarea, #${elemId} input`);
if(!f){ boundField=null; usingSetterHook=false; }
else if(f!==boundField){ hookField(f); }
setTimeout(pollField, 300);
}
function pollValue(){
if(usingSetterHook){ setTimeout(pollValue, 500); return; }
const f = document.querySelector(`#${elemId} textarea, #${elemId} input`);
if(!f){ setTimeout(pollValue,300); return; }
const v = f.value;
if(v&&v!==lastValue){ lastValue=v; handlePayload(v); }
setTimeout(pollValue, 40);
}
function tryUnlockAudio(){
if(!audioCtx) audioCtx = new (window.AudioContext||window.webkitAudioContext)();
if(audioCtx.state==="suspended") audioCtx.resume();
}
document.addEventListener("click", e => { if(e.target.closest("#tts_generate")) tryUnlockAudio(); });
pollField(); pollValue();
"""
with gr.Blocks(title="MOSS-TTS-Realtime") as demo:
gr.Markdown(f"""
# 🎙️ MOSS-TTS-Realtime
**Context-aware, multi-turn streaming TTS with real-time Web Audio playback.**
**Supported languages:** {SUPPORTED_LANGUAGES}
> 💡 Upload a 10–30s reference audio clip for voice cloning.
> Without reference audio the model uses a random speaker voice.
""")
with gr.Row():
with gr.Column(scale=3):
user_text = gr.Textbox(
label="User Text (conversation context, optional)",
placeholder="e.g. Hello, how are you?", lines=2)
assistant_text = gr.Textbox(
label="Assistant Text (what to synthesize) ✱ required",
placeholder="Enter the text you want to synthesize here…", lines=5)
prompt_audio = gr.Audio(
label="🎵 Reference Audio (optional — voice cloning)",
type="filepath", sources=["upload", "microphone"])
gr.HTML(STREAM_PLAYER_HTML, js_on_load=STREAM_PLAYER_JS)
run_btn = gr.Button("▶ Generate", variant="primary",
elem_id="tts_generate")
stream_data = gr.Textbox(label="PCM Stream", elem_id="pcm_stream",
interactive=False, lines=2)
output_audio = gr.Audio(label="🔊 Final Audio", type="numpy")
status = gr.Textbox(label="📊 Status / Stats",
lines=2, interactive=False)
with gr.Column(scale=2):
gr.Markdown("### ⚙️ Generation Parameters")
temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.05,
label="Temperature")
top_p = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="Top-P")
top_k = gr.Slider(1, 100, value=30, step=1, label="Top-K")
repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.05,
label="Repetition Penalty")
repetition_window = gr.Slider(1, 200, value=50, step=1,
label="Repetition Window")
gr.Markdown("### 🔧 Streaming Options")
text_chunk_tokens = gr.Slider(1, 64, value=12, step=1,
label="Text Chunk Tokens",
info="Tokens pushed to model per step — smaller = lower TTFB")
decode_chunk_frames = gr.Slider(0, 20, value=12, step=1,
label="Decode Chunk Frames")
decode_overlap_frames = gr.Slider(0, 10, value=0, step=1,
label="Decode Overlap Frames")
chunk_duration = gr.Slider(0.0, 1.0, value=0.24, step=0.01,
label="Codec Chunk Duration (s)")
gr.Markdown("""
---
### ℹ️ Notes
- **TTFB** = time until the first audio chunk arrives in the browser.
- **RTF < 1.0** = faster than real-time generation.
- Audio plays live in the browser via Web Audio API as chunks arrive.
- The final waveform is also shown in the audio player below the status box.
""")
run_btn.click(
fn=generate,
inputs=[user_text, assistant_text, prompt_audio,
temperature, top_p, top_k,
repetition_penalty, repetition_window,
text_chunk_tokens,
decode_chunk_frames, decode_overlap_frames, chunk_duration],
outputs=[stream_data, output_audio, status],
)
if __name__ == "__main__":
demo.queue(max_size=8, default_concurrency_limit=1).launch(
server_name="0.0.0.0",
server_port=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),
ssr_mode=False,
theme=gr.themes.Soft(),
)