| import os |
| import tempfile |
| import threading |
| import time |
| import wave |
| from pathlib import Path |
|
|
| os.environ.setdefault("HF_HOME", "/tmp/.cache/huggingface") |
| os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules") |
| os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") |
| os.environ.setdefault("ZONOS2_TTS_NORM_CACHE_DIR", "/tmp/zonos2-tts-norm") |
| os.environ.setdefault("GRADIO_SSR_MODE", "false") |
| os.environ.setdefault("NUMBA_DISABLE_CUDA", "1") |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
|
|
| for cache_dir in ( |
| os.environ["HF_HOME"], |
| os.environ["HF_MODULES_CACHE"], |
| os.environ["MPLCONFIGDIR"], |
| os.environ["ZONOS2_TTS_NORM_CACHE_DIR"], |
| ): |
| Path(cache_dir).mkdir(parents=True, exist_ok=True) |
|
|
|
|
| print("Importing Space runtime dependencies...", flush=True) |
| import spaces |
| import gradio as gr |
| import numpy as np |
| import torch |
|
|
| print("Importing ZONOS2 modules...", flush=True) |
| from zonos2.message import TTSSamplingParams |
| from zonos2.tokenizer.textnorm import SERVER_TO_NEMO_LANG, TTSTextNormalizer |
| from zonos2.tts import TTSLLM |
| print("Imported ZONOS2 modules.", flush=True) |
|
|
| MODEL_ID = "Zyphra/ZONOS2" |
| SAMPLE_RATE = 44100 |
| LANGUAGES = [ |
| ("English (US)", "en_us"), |
| ("English (UK)", "en_gb"), |
| ("French", "fr_fr"), |
| ("German", "de"), |
| ("Spanish", "es"), |
| ("Italian", "it"), |
| ("Portuguese (Brazil)", "pt_br"), |
| ("Japanese", "ja"), |
| ("Mandarin Chinese", "cmn"), |
| ("Korean", "ko"), |
| ] |
| SPEAKING_RATE_BUCKETS = [ |
| ("Default", "default"), |
| ("Very slow", "0"), |
| ("Slow", "1"), |
| ("Relaxed", "2"), |
| ("Natural", "3"), |
| ("Bright", "4"), |
| ("Fast", "5"), |
| ("Very fast", "6"), |
| ("Extreme", "7"), |
| ] |
|
|
| torch.backends.cuda.matmul.allow_tf32 = True |
|
|
|
|
| def _load_model() -> TTSLLM: |
| print(f"Loading {MODEL_ID} for ZeroGPU inference...", flush=True) |
| started = time.perf_counter() |
| model = TTSLLM( |
| model_path=MODEL_ID, |
| decode_audio=True, |
| cuda_graph_max_bs=0, |
| max_running_req=4, |
| max_extend_tokens=4096, |
| memory_ratio=0.75, |
| use_pynccl=False, |
| ) |
| elapsed = time.perf_counter() - started |
| print(f"Loaded {MODEL_ID} in {elapsed:.1f}s", flush=True) |
| return model |
|
|
|
|
| TTS: TTSLLM | None = None |
| TEXT_NORMALIZER = TTSTextNormalizer() |
| TTS_LOCK = threading.Lock() |
|
|
|
|
| def _estimate_duration(*args, **kwargs) -> int: |
| max_tokens = kwargs.get("max_tokens") |
| if max_tokens is None and len(args) > 4: |
| max_tokens = args[4] |
| try: |
| max_tokens = int(max_tokens) |
| except (TypeError, ValueError): |
| max_tokens = 768 |
| base_seconds = 220 if TTS is None else 45 |
| return min(300, max(60, base_seconds + max_tokens // 12)) |
|
|
|
|
| def _pcm_float32_to_wav(audio_bytes: bytes, sample_rate: int = SAMPLE_RATE) -> str: |
| audio = np.frombuffer(audio_bytes, dtype=np.float32) |
| if audio.size == 0: |
| raise gr.Error("The model returned no audio. Try increasing max tokens.") |
| audio = np.nan_to_num(audio, nan=0.0, posinf=0.0, neginf=0.0) |
| audio = np.clip(audio, -1.0, 1.0) |
| audio_i16 = (audio * 32767.0).astype(np.int16) |
|
|
| handle = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) |
| handle.close() |
| with wave.open(handle.name, "wb") as wav: |
| wav.setnchannels(1) |
| wav.setsampwidth(2) |
| wav.setframerate(sample_rate) |
| wav.writeframes(audio_i16.tobytes()) |
| return handle.name |
|
|
|
|
| def _normalize_text(text: str, language: str, enabled: bool) -> str: |
| if not enabled: |
| return text |
| if language not in SERVER_TO_NEMO_LANG: |
| return text |
| return TEXT_NORMALIZER.normalize(text, language) |
|
|
|
|
| def _speaking_rate_bucket(value: str) -> int | None: |
| if value == "default": |
| return None |
| return int(value) |
|
|
|
|
| @spaces.GPU(duration=_estimate_duration) |
| def synthesize( |
| text: str, |
| language: str, |
| text_normalization: bool, |
| speaking_rate: str, |
| max_tokens: int, |
| temperature: float, |
| topk: int, |
| top_p: float, |
| min_p: float, |
| repetition_penalty: float, |
| seed: int, |
| ): |
| text = (text or "").strip() |
| if not text: |
| raise gr.Error("Enter text to synthesize.") |
| if len(text) > 1200: |
| raise gr.Error("Keep the prompt under 1200 characters for this Space.") |
|
|
| normalized = _normalize_text(text, language, text_normalization) |
| params = TTSSamplingParams( |
| temperature=float(temperature), |
| topk=int(topk), |
| top_p=float(top_p), |
| min_p=float(min_p), |
| max_tokens=int(max_tokens), |
| repetition_window=50, |
| repetition_penalty=float(repetition_penalty), |
| repetition_codebooks=8, |
| seed=None if seed is None or int(seed) < 0 else int(seed), |
| ) |
|
|
| started = time.perf_counter() |
| with TTS_LOCK: |
| global TTS |
| if TTS is None: |
| TTS = _load_model() |
| torch.cuda.set_stream(TTS.stream) |
| result = TTS.generate_one( |
| normalized, |
| params, |
| decode_audio=True, |
| speaking_rate_bucket=_speaking_rate_bucket(speaking_rate), |
| quality_buckets=None, |
| ) |
| elapsed = time.perf_counter() - started |
|
|
| wav_path = _pcm_float32_to_wav(result["audio"], result.get("sample_rate", SAMPLE_RATE)) |
| frames = len(result.get("audio_tokens") or []) |
| eos_frame = result.get("eos_frame") |
| status = f"Generated {frames} frames in {elapsed:.1f}s" |
| if eos_frame is not None: |
| status += f" (EOS frame {eos_frame})" |
| if normalized != text: |
| status += f"\n\nNormalized text: {normalized}" |
| return wav_path, status |
|
|
|
|
| CSS = """ |
| main, .gradio-container, .gradio-container > .fillable { |
| max-width: 1180px !important; |
| margin-inline: auto !important; |
| } |
| .compact-status textarea { |
| font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; |
| } |
| """ |
|
|
|
|
| with gr.Blocks(title="ZONOS2") as demo: |
| gr.Markdown("# ZONOS2") |
| with gr.Row(): |
| with gr.Column(scale=5): |
| text = gr.Textbox( |
| label="Text", |
| value="In the quiet hum of the studio, ZONOS2 turns written words into natural speech.", |
| lines=6, |
| max_length=1200, |
| ) |
| with gr.Row(): |
| language = gr.Dropdown( |
| choices=LANGUAGES, |
| value="en_us", |
| label="Language", |
| ) |
| speaking_rate = gr.Dropdown( |
| choices=SPEAKING_RATE_BUCKETS, |
| value="default", |
| label="Speaking rate", |
| ) |
| text_normalization = gr.Checkbox(value=True, label="Text normalization") |
| generate = gr.Button("Generate", variant="primary") |
| with gr.Column(scale=4): |
| audio = gr.Audio(label="Audio", type="filepath", format="wav") |
| status = gr.Textbox( |
| label="Status", |
| lines=5, |
| interactive=False, |
| elem_classes=["compact-status"], |
| ) |
|
|
| with gr.Accordion("Sampling", open=False): |
| with gr.Row(): |
| max_tokens = gr.Slider( |
| minimum=128, |
| maximum=2048, |
| step=64, |
| value=768, |
| label="Max audio tokens", |
| ) |
| seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)") |
| with gr.Row(): |
| temperature = gr.Slider(0.1, 2.0, value=1.15, step=0.05, label="Temperature") |
| topk = gr.Slider(1, 512, value=106, step=1, label="Top-k") |
| with gr.Row(): |
| top_p = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Top-p") |
| min_p = gr.Slider(0.0, 0.5, value=0.18, step=0.01, label="Min-p") |
| repetition_penalty = gr.Slider( |
| 1.0, |
| 2.0, |
| value=1.2, |
| step=0.05, |
| label="Repetition penalty", |
| ) |
|
|
| gr.Examples( |
| examples=[ |
| [ |
| "The first explorers landed just after sunrise, carrying maps, coffee, and impossible optimism.", |
| "en_us", |
| True, |
| "default", |
| 512, |
| 1.15, |
| 106, |
| 0.0, |
| 0.18, |
| 1.2, |
| -1, |
| ], |
| [ |
| "Le modèle parle avec une voix claire, expressive et naturellement rythmée.", |
| "fr_fr", |
| True, |
| "default", |
| 512, |
| 1.15, |
| 106, |
| 0.0, |
| 0.18, |
| 1.2, |
| -1, |
| ], |
| ], |
| inputs=[ |
| text, |
| language, |
| text_normalization, |
| speaking_rate, |
| max_tokens, |
| temperature, |
| topk, |
| top_p, |
| min_p, |
| repetition_penalty, |
| seed, |
| ], |
| ) |
|
|
| generate.click( |
| fn=synthesize, |
| inputs=[ |
| text, |
| language, |
| text_normalization, |
| speaking_rate, |
| max_tokens, |
| temperature, |
| topk, |
| top_p, |
| min_p, |
| repetition_penalty, |
| seed, |
| ], |
| outputs=[audio, status], |
| api_name="generate", |
| concurrency_limit=1, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.queue(max_size=8, default_concurrency_limit=1).launch(css=CSS) |
|
|
