Upload /tmp/hugging-demos-build-model_ZDisket_MOSS-TTS-PNY-bxj3g8d1/portable_tts_runtime.py with huggingface_hub

tmp/hugging-demos-build-model_ZDisket_MOSS-TTS-PNY-bxj3g8d1/portable_tts_runtime.py

@@ -0,0 +1,462 @@
+"""Gradio demo for ZDisket/MOSS-TTS-PNY — a speaker-conditioned TTS model.
+
+This Space loads the MOSS-TTS-PNY checkpoint (fine-tuned on MLP:FiM and TF2
+voices) and provides a Gradio interface for text-to-speech synthesis with
+speaker selection, emotion control, and energy adjustment.
+
+Model: https://huggingface.co/ZDisket/MOSS-TTS-PNY
+"""
+
+from __future__ import annotations
+
+import csv
+import json
+import math
+import os
+import random
+import tempfile
+import time
+import unicodedata
+import uuid
+from pathlib import Path
+from typing import Any
+
+# ── ZeroGPU: import spaces before any CUDA-touching import ──────────────
+import spaces  # noqa: E402
+
+import torch  # noqa: E402
+import gradio as gr  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+
+# ── Path setup ───────────────────────────────────────────────────────────
+BUILD_ROOT = Path(__file__).resolve().parent
+MODEL_REPO = "ZDisket/MOSS-TTS-PNY"
+
+# Skip the os.execv re-exec inside portable_tts_runtime
+os.environ.setdefault("MOSS_TTS_NVIDIA_LD_LIBRARY_PATH_READY", "1")
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+# ── Download model weights at startup ─────────────────────────────────────
+print("Downloading model weights from ZDisket/MOSS-TTS-PNY …", flush=True)
+_download_start = time.perf_counter()
+_model_root = snapshot_download(
+    repo_id=MODEL_REPO,
+    repo_type="model",
+    allow_patterns=[
+        "moss_tts_local_clipper_checkpoint/**",
+        "moss_audio_tokenizer/**",
+        "istftnet2_decoder4_50hz/**",
+    ],
+)
+_download_elapsed = time.perf_counter() - _download_start
+print(f"Model weights downloaded in {_download_elapsed:.1f}s -> {_model_root}", flush=True)
+
+CHECKPOINT_PATH = str(Path(_model_root) / "moss_tts_local_clipper_checkpoint")
+CODEC_PATH = str(Path(_model_root) / "moss_audio_tokenizer")
+DECODER_DIR = str(Path(_model_root) / "istftnet2_decoder4_50hz")
+
+# ── Constants ─────────────────────────────────────────────────────────────
+CODEC_FRAMES_PER_SECOND = 12.5
+MAX_OUTPUT_SECONDS = 30
+MAX_OUTPUT_FRAMES = int(CODEC_FRAMES_PER_SECOND * MAX_OUTPUT_SECONDS)
+MAX_TEXT_LINES = 12
+MAX_LANGUAGE_CHARS = 16
+MAX_TEXT_CHARS = 500
+
+EMOTION_CHOICES = [
+    ("Neutral", 0),
+    ("Happy", 1),
+    ("Sad", 2),
+    ("Angry", 3),
+    ("Annoyed", 4),
+    ("Fearful", 5),
+    ("Surprised", 6),
+    ("Calm", 7),
+    ("Disgusted", 8),
+    ("Whispering", 9),
+    ("Nonverbal", 10),
+    ("Shouting", 11),
+]
+
+RANDOM_TEXTS = [
+    "I know this sounds sudden, but I really missed hearing your voice today.",
+    "Wait, you actually finished the whole thing already?",
+    "I'm trying to stay calm, but that was way closer than I expected.",
+    "Honestly, I'm proud of us for getting through that without giving up.",
+    "No, no, it's fine. I'm only a little dramatically offended.",
+    "I cannot believe you kept that surprise hidden for so long.",
+    "Please tell me you saw that too, because I am not ready to explain it.",
+    "That was kind of beautiful, in a messy and completely unexpected way.",
+    "I'm sorry, I should have listened before jumping to conclusions.",
+    "Okay, deep breath. We can fix this if we take it one step at a time.",
+    "You have no idea how happy that made me.",
+    "I was nervous at first, but now I'm actually excited to try again.",
+]
+
+# ── Load speaker map & tiers ──────────────────────────────────────────────
+
+
+def load_speaker_tiers(csv_path: str) -> dict[int, int]:
+    path = Path(csv_path)
+    if not path.exists():
+        return {}
+    tiers: dict[int, int] = {}
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        for row in csv.DictReader(handle):
+            tiers[int(row["speaker_id"])] = int(row["tier"])
+    return tiers
+
+
+def load_speaker_choices(csv_path: str, tiers: dict[int, int]) -> tuple[list[str], dict[str, int]]:
+    path = Path(csv_path)
+    if not path.exists():
+        choices = [f"Speaker {i} ({i})" for i in range(43)]
+        return choices, {c: i for i, c in enumerate(choices)}
+    choice_rows: list[tuple[int, str, str, int]] = []
+    with path.open("r", encoding="utf-8", newline="") as handle:
+        for row in csv.DictReader(handle):
+            speaker_id = int(row["speaker_id"])
+            label = row.get("label") or f"{row.get('display_name') or row.get('name')} ({speaker_id})"
+            tier = tiers.get(speaker_id, 9)
+            if tier != 9:
+                label = f"T{tier} - {label}"
+            choice_rows.append((tier, label.lower(), label, speaker_id))
+    choices: list[str] = []
+    speaker_ids: dict[str, int] = {}
+    for _tier, _sort_label, label, speaker_id in sorted(choice_rows):
+        choices.append(label)
+        speaker_ids[label] = speaker_id
+    return choices, speaker_ids
+
+
+SPEAKER_TIERS = load_speaker_tiers(str(BUILD_ROOT / "speaker_hours_tiers.csv"))
+SPEAKER_CHOICES, SPEAKER_IDS = load_speaker_choices(
+    str(BUILD_ROOT / "speaker_id_map.csv"), SPEAKER_TIERS
+)
+DEFAULT_SPEAKER = next(
+    (c for c in SPEAKER_CHOICES if "Twilight (31)" in c),
+    SPEAKER_CHOICES[0] if SPEAKER_CHOICES else "Speaker 31 (31)",
+)
+
+# ── Load model at module scope (ZeroGPU rule 2) ───────────────────────────
+from portable_tts_runtime import (  # noqa: E402
+    OptimizedTTSConfig,
+    OptimizedTTSRunner,
+    save_wav_pcm16,
+)
+
+# Use fixed-full torch opt mode (avoids the ~4-minute Triton compile of the
+# packed-local path; still gets torch.compile on the global transformer).
+RUNNER = OptimizedTTSRunner(
+    OptimizedTTSConfig(
+        checkpoint=CHECKPOINT_PATH,
+        codec_path=CODEC_PATH,
+        decoder_dir=DECODER_DIR,
+        decoder_runtime="torchscript_cuda",
+        decoder4_features_runtime="torch_fp16",
+        dtype="fp16",
+        attn_implementation="sdpa",
+        torch_opt_mode="fixed-full",
+        cache_implementation="static",
+        compile_global_transformer=True,
+        global_compile_mode="default",
+        fast_prepare_inputs=True,
+        triton_top_p=True,
+        packed_local_qkv=False,
+        packed_local_mlp=False,
+        static_packed_weights=False,
+        tensorize_rmsnorm_eps=True,
+        feedback_lookup=True,
+        vocoder_bucket_frames=64,
+        vocoder_prewarm_buckets="64,128,192,256,320,384,448,512,576,640",
+    )
+)
+SAMPLE_RATE = int(RUNNER.sample_rate)
+print(f"Model loaded. sample_rate={SAMPLE_RATE}", flush=True)
+
+# NOTE: prewarm is skipped at module scope on ZeroGPU — no real GPU is
+# attached to the main process. The first @spaces.GPU call will pay the
+# compile/warmup cost.
+
+
+# ── Helpers ────────────────────────────────────────────────────────────────
+
+
+def energy_label(energy: float) -> str:
+    if energy < 0.33:
+        return "low energy"
+    if energy < 0.66:
+        return "medium energy"
+    return "high energy"
+
+
+def normalize_client_text(value: str, *, max_chars: int) -> str:
+    if not isinstance(value, str):
+        raise gr.Error("Text must be a string.")
+    text = unicodedata.normalize("NFKC", value).replace("\r\n", "\n").replace("\r", "\n").strip()
+    if not text:
+        raise gr.Error("Enter some text first.")
+    if max_chars > 0 and len(text) > max_chars:
+        raise gr.Error(f"Text is too long. Keep it under {max_chars} characters.")
+    if text.count("\n") + 1 > MAX_TEXT_LINES:
+        raise gr.Error(f"Text has too many lines. Keep it under {MAX_TEXT_LINES} lines.")
+    for char in text:
+        cat = unicodedata.category(char)
+        if cat.startswith("C") and char not in {"\n", "\t"}:
+            raise gr.Error("Text contains unsupported control characters.")
+    if not any(c.isalnum() for c in text):
+        raise gr.Error("Text must contain at least one letter or number.")
+    if "\ufffd" in text:
+        raise gr.Error("Text contains malformed replacement characters.")
+    return text
+
+
+def normalize_language(value: str) -> str:
+    language = unicodedata.normalize("NFKC", value or "en").strip().lower()
+    if not language:
+        return "en"
+    if len(language) > MAX_LANGUAGE_CHARS:
+        raise gr.Error("Language code is too long.")
+    allowed = set("abcdefghijklmnopqrstuvwxyz-")
+    if any(c not in allowed for c in language):
+        raise gr.Error("Language must be a simple code like 'en'.")
+    return language
+
+
+def validate_float(name: str, value: float, minimum: float, maximum: float) -> float:
+    value = float(value)
+    if not math.isfinite(value) or value < minimum or value > maximum:
+        raise gr.Error(f"{name} must be between {minimum} and {maximum}.")
+    return value
+
+
+def validate_int(name: str, value: int, minimum: int, maximum: int) -> int:
+    value = int(value)
+    if value < minimum or value > maximum:
+        raise gr.Error(f"{name} must be between {minimum} and {maximum}.")
+    return value
+
+
+def _random_text() -> str:
+    return random.choice(RANDOM_TEXTS)
+
+
+# ── Inference (decorated for ZeroGPU) ─────────────────────────────────────
+
+
+@spaces.GPU(duration=120)
+def synthesize(
+    text: str,
+    style_text: str,
+    language: str,
+    speaker_label: str,
+    emotion_label: str,
+    style_energy: float,
+    audio_temperature: float,
+    audio_top_p: float,
+    audio_top_k: int,
+    audio_repetition_penalty: float,
+) -> tuple[str, str]:
+    """Generate speech audio from text using MOSS-TTS-PNY.
+
+    Args:
+        text: The text to synthesize.
+        style_text: Optional style reference text (falls back to ``text``).
+        language: Language code (e.g. ``"en"``).
+        speaker_label: Speaker name from the dropdown.
+        emotion_label: Emotion name from the dropdown.
+        style_energy: Energy level (0.0 – 1.0).
+        audio_temperature: Sampling temperature for audio tokens.
+        audio_top_p: Top-p value for audio token sampling.
+        audio_top_k: Top-k value (0 = disabled).
+        audio_repetition_penalty: Repetition penalty for audio tokens.
+
+    Returns:
+        A tuple of (wav_filepath, stats_json).
+    """
+    clean_text = normalize_client_text(text, max_chars=MAX_TEXT_CHARS)
+    clean_style_text = (
+        normalize_client_text(style_text, max_chars=MAX_TEXT_CHARS)
+        if style_text and style_text.strip()
+        else ""
+    )
+    clean_language = normalize_language(language)
+    if speaker_label not in SPEAKER_IDS:
+        raise gr.Error("Choose a valid speaker.")
+    emotion_lookup = {label: eid for label, eid in EMOTION_CHOICES}
+    if emotion_label not in emotion_lookup:
+        raise gr.Error("Choose a valid emotion.")
+    clean_energy = validate_float("Energy", style_energy, 0.0, 1.0)
+    clean_temperature = validate_float("Audio temperature", audio_temperature, 0.0, 1.5)
+    clean_top_p = validate_float("Audio top-p", audio_top_p, 0.05, 1.0)
+    clean_top_k = validate_int("Audio top-k", audio_top_k, 0, 200)
+    clean_rep_pen = validate_float(
+        "Audio repetition penalty", audio_repetition_penalty, 0.0, 2.0
+    )
+
+    start = time.perf_counter()
+    result = RUNNER.synthesize(
+        text=clean_text,
+        language=clean_language,
+        speaker_id=SPEAKER_IDS[speaker_label],
+        max_new_tokens=MAX_OUTPUT_FRAMES,
+        text_temperature=0.0,
+        text_top_p=1.0,
+        text_top_k=None,
+        audio_temperature=clean_temperature,
+        audio_top_p=clean_top_p,
+        audio_top_k=clean_top_k if clean_top_k > 0 else None,
+        audio_repetition_penalty=clean_rep_pen,
+        n_vq_for_inference=32,
+        style_text=clean_style_text or clean_text,
+        style_emotion_id=emotion_lookup[emotion_label],
+        style_energy=clean_energy,
+    )
+    elapsed = time.perf_counter() - start
+
+    # Write to a temp file (concurrency-safe — no fixed output paths).
+    tmp_wav = tempfile.NamedTemporaryFile(
+        suffix=".wav", delete=False, dir=str(BUILD_ROOT / "outputs")
+    )
+    tmp_wav.close()
+    output_path = Path(tmp_wav.name)
+    save_wav_pcm16(output_path, result["audio"], SAMPLE_RATE)
+
+    stats = {
+        "speaker": speaker_label,
+        "speaker_id": SPEAKER_IDS[speaker_label],
+        "emotion": emotion_label,
+        "style_energy": clean_energy,
+        "energy_label": energy_label(clean_energy),
+        "language": clean_language,
+        "text_chars": len(clean_text),
+        "audio_sec": round(result["audio_sec"], 3),
+        "prompt_sec": round(result["prompt_sec"], 3),
+        "generate_sec": round(result["generate_sec"], 3),
+        "decode_sec": round(result["decode_sec"], 3),
+        "total_sec": round(elapsed, 3),
+        "generate_x_realtime": round(result["generate_x_realtime"], 2),
+        "generated_tokens": result.get("generated_tokens", 0),
+        "prompt_tokens": result.get("prompt_tokens", 0),
+        "sample_rate": SAMPLE_RATE,
+    }
+    return str(output_path), json.dumps(stats, indent=2)
+
+
+# ── UI ────────────────────────────────────────────────────────────────────
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(
+    title="MOSS-TTS-PNY",
+    theme=gr.themes.Citrus(),
+    css=CSS,
+) as demo:
+    gr.Markdown(
+        "# 🎙️ MOSS-TTS-PNY\n"
+        "Speaker-conditioned text-to-speech with emotion and energy control. "
+        "Fine-tuned MOSS-TTS checkpoint featuring character voices.\n\n"
+        "Model: [ZDisket/MOSS-TTS-PNY](https://huggingface.co/ZDisket/MOSS-TTS-PNY)"
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            text = gr.Textbox(
+                label="Text",
+                value="I wasn't expecting that to work, but now I'm kind of excited.",
+                lines=4,
+                max_length=MAX_TEXT_CHARS,
+            )
+            style_text = gr.Textbox(
+                label="Style text (optional)",
+                placeholder="Leave empty to use the same text as input.",
+                lines=2,
+                max_length=MAX_TEXT_CHARS,
+            )
+            with gr.Row():
+                randomize = gr.Button("Random Text")
+                generate = gr.Button("Generate", variant="primary")
+
+        with gr.Column(scale=2):
+            language = gr.Textbox(label="Language", value="en")
+            speaker = gr.Dropdown(
+                label="Speaker",
+                choices=SPEAKER_CHOICES,
+                value=DEFAULT_SPEAKER,
+            )
+            emotion = gr.Dropdown(
+                label="Emotion",
+                choices=[label for label, _ in EMOTION_CHOICES],
+                value="Neutral",
+            )
+            style_energy = gr.Slider(
+                label="Energy",
+                minimum=0.0,
+                maximum=1.0,
+                value=0.5,
+                step=0.05,
+            )
+            gr.Markdown(f"Max {MAX_OUTPUT_SECONDS}s of audio output")
+
+    with gr.Accordion("Advanced settings", open=False):
+        gr.Markdown("If you don't know what you're doing, leave these untouched.")
+        with gr.Row():
+            audio_temperature = gr.Slider(
+                label="Audio temperature", minimum=0.0, maximum=1.5, value=0.8, step=0.05
+            )
+            audio_top_p = gr.Slider(
+                label="Audio top-p", minimum=0.05, maximum=1.0, value=0.92, step=0.05
+            )
+            audio_top_k = gr.Slider(
+                label="Audio top-k", minimum=0, maximum=200, value=0, step=1
+            )
+            audio_repetition_penalty = gr.Slider(
+                label="Audio repetition penalty",
+                minimum=0.0,
+                maximum=2.0,
+                value=1.05,
+                step=0.05,
+            )
+
+    with gr.Row():
+        audio_output = gr.Audio(label="Generated audio", type="filepath", autoplay=True)
+        stats = gr.Code(label="Stats", language="json")
+
+    gr.Examples(
+        examples=[
+            ["I wasn't expecting that to work, but now I'm kind of excited.", "", "en",
+             DEFAULT_SPEAKER, "Neutral", 0.5, 0.8, 0.92, 0, 1.05],
+            ["You have no idea how happy that made me.", "", "en",
+             DEFAULT_SPEAKER, "Happy", 0.8, 0.8, 0.92, 0, 1.05],
+            ["I'm trying to stay calm, but that was way closer than I expected.", "", "en",
+             DEFAULT_SPEAKER, "Fearful", 0.3, 0.8, 0.92, 0, 1.05],
+            ["Okay, deep breath. We can fix this if we take it one step at a time.", "", "en",
+             DEFAULT_SPEAKER, "Calm", 0.4, 0.8, 0.92, 0, 1.05],
+        ],
+        inputs=[
+            text, style_text, language, speaker, emotion, style_energy,
+            audio_temperature, audio_top_p, audio_top_k, audio_repetition_penalty,
+        ],
+        outputs=[audio_output, stats],
+        fn=synthesize,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    randomize.click(_random_text, inputs=[], outputs=[text], queue=False)
+    generate.click(
+        synthesize,
+        inputs=[
+            text, style_text, language, speaker, emotion, style_energy,
+            audio_temperature, audio_top_p, audio_top_k, audio_repetition_penalty,
+        ],
+        outputs=[audio_output, stats],
+        api_name="generate",
+    )
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1, max_size=20)
+    demo.launch(mcp_server=True, show_error=False)