OpenMOSS-Team/MOSS-TTS-v1.5

1import argparse2import functools3import importlib.util4import inspect5import os6from pathlib import Path7import re8import time9import orjson10 11try:12    import spaces13except ImportError:14    class _SpacesFallback:15        @staticmethod16        def GPU(*_args, **_kwargs):17            def _decorator(func):18                return func19 20            return _decorator21 22    spaces = _SpacesFallback()23 24import gradio as gr25import numpy as np26import torch27from transformers import AutoModel, AutoProcessor28 29# Disable the broken cuDNN SDPA backend30torch.backends.cuda.enable_cudnn_sdp(False)31# Keep these enabled as fallbacks32torch.backends.cuda.enable_flash_sdp(True)33torch.backends.cuda.enable_mem_efficient_sdp(True)34torch.backends.cuda.enable_math_sdp(True)35 36MODEL_PATH = "OpenMOSS-Team/MOSS-TTS-v1.5"37DEFAULT_ATTN_IMPLEMENTATION = "auto"38DEFAULT_MAX_NEW_TOKENS = 409639PRELOAD_ENV_VAR = "MOSS_TTS_PRELOAD_AT_STARTUP"40CONTINUATION_NOTICE = (41    "Continuation mode is active. Make sure the reference audio transcript is prepended to the input text."42)43 44MODE_CLONE = "Clone"45MODE_CONTINUE = "Continuation"46MODE_CONTINUE_CLONE = "Continuation + Clone"47ZH_TOKENS_PER_CHAR = 3.09841195131303348EN_TOKENS_PER_CHAR = 0.867337626275521949REFERENCE_AUDIO_DIR = Path(__file__).resolve().parent / "assets" / "audio"50EXAMPLE_TEXTS_JSONL_PATH = Path(__file__).resolve().parent / "assets" / "text" / "moss_tts_example_texts.jsonl"51LANGUAGE_TAG_AUTO = "Auto (omit)"52LANGUAGE_TAG_CHOICES = [53    LANGUAGE_TAG_AUTO,54    "Chinese",55    "Cantonese",56    "English",57    "Arabic",58    "Czech",59    "Danish",60    "Dutch",61    "Finnish",62    "French",63    "German",64    "Greek",65    "Hebrew",66    "Hindi",67    "Hungarian",68    "Italian",69    "Japanese",70    "Korean",71    "Macedonian",72    "Malay",73    "Persian (Farsi)",74    "Polish",75    "Portuguese",76    "Romanian",77    "Russian",78    "Spanish",79    "Swahili",80    "Swedish",81    "Tagalog",82    "Thai",83    "Turkish",84    "Vietnamese",85]86 87 88def _parse_example_id(example_id: str) -> tuple[str, int] | None:89    matched = re.fullmatch(r"(zh|en)/(\d+)", (example_id or "").strip())90    if matched is None:91        return None92    return matched.group(1), int(matched.group(2))93 94 95def _resolve_reference_audio_path(language: str, index: int) -> Path | None:96    stem_candidates = [f"reference_{language}_{index}"]97    for stem in stem_candidates:98        for ext in (".wav", ".mp3"):99            audio_path = REFERENCE_AUDIO_DIR / f"{stem}{ext}"100            if audio_path.exists():101                return audio_path102    return None103 104 105def build_example_rows() -> list[tuple[str, str, str]]:106    rows: list[tuple[str, str, str]] = []107 108    with open(EXAMPLE_TEXTS_JSONL_PATH, "rb") as f:109        for line in f:110            if not line.strip():111                continue112            sample = orjson.loads(line)113            parsed = _parse_example_id(sample.get("id", ""))114            if parsed is None:115                continue116 117            language, index = parsed118            text = str(sample.get("text", "")).strip()119            audio_path = _resolve_reference_audio_path(language, index)120            if audio_path is None:121                continue122 123            rows.append((sample['role'], str(audio_path), text))124 125    return rows126 127 128EXAMPLE_ROWS = build_example_rows()129 130 131@functools.lru_cache(maxsize=1)132def load_backend(model_path: str, device_str: str, attn_implementation: str):133    device = torch.device(device_str if torch.cuda.is_available() else "cpu")134    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32135    resolved_attn_implementation = resolve_attn_implementation(136        requested=attn_implementation,137        device=device,138        dtype=dtype,139    )140 141    processor = AutoProcessor.from_pretrained(142        model_path,143        trust_remote_code=True,144    )145    if hasattr(processor, "audio_tokenizer"):146        processor.audio_tokenizer = processor.audio_tokenizer.to(device)147        processor.audio_tokenizer.eval()148 149    model_kwargs = {150        "trust_remote_code": True,151        "torch_dtype": dtype,152    }153    if resolved_attn_implementation:154        model_kwargs["attn_implementation"] = resolved_attn_implementation155 156    model = AutoModel.from_pretrained(model_path, **model_kwargs).to(device)157    model.eval()158 159    sample_rate = int(getattr(processor.model_config, "sampling_rate", 24000))160    return model, processor, device, sample_rate161 162 163def resolve_attn_implementation(requested: str, device: torch.device, dtype: torch.dtype) -> str | None:164    requested_norm = (requested or "").strip().lower()165 166    if requested_norm in {"none"}:167        return None168 169    if requested_norm not in {"", "auto"}:170        return requested171 172    # Prefer FlashAttention 2 when package + device conditions are met.173    if (174        device.type == "cuda"175        and importlib.util.find_spec("flash_attn") is not None176        and dtype in {torch.float16, torch.bfloat16}177    ):178        major, _ = torch.cuda.get_device_capability(device)179        if major >= 8:180            return "flash_attention_2"181 182    # CUDA fallback: use PyTorch SDPA kernels.183    if device.type == "cuda":184        return "sdpa"185 186    # CPU fallback.187    return "eager"188 189 190def detect_text_language(text: str) -> str:191    zh_chars = len(re.findall(r"[\u4e00-\u9fff]", text))192    en_chars = len(re.findall(r"[A-Za-z]", text))193    if zh_chars == 0 and en_chars == 0:194        return "en"195    return "zh" if zh_chars >= en_chars else "en"196 197 198def supports_duration_control(mode_with_reference: str) -> bool:199    return mode_with_reference not in {MODE_CONTINUE, MODE_CONTINUE_CLONE}200 201 202def estimate_duration_tokens(text: str) -> tuple[str, int, int, int]:203    normalized = text or ""204    effective_len = max(len(normalized), 1)205    language = detect_text_language(normalized)206    factor = ZH_TOKENS_PER_CHAR if language == "zh" else EN_TOKENS_PER_CHAR207    default_tokens = max(1, int(effective_len * factor))208    min_tokens = max(1, int(default_tokens * 0.5))209    max_tokens = max(min_tokens, int(default_tokens * 1.5))210    return language, default_tokens, min_tokens, max_tokens211 212 213def update_duration_controls(214    enabled: bool,215    text: str,216    current_tokens: float | int | None,217    mode_with_reference: str,218):219    if not supports_duration_control(mode_with_reference):220        return (221            gr.update(visible=False),222            "Duration control is disabled for Continuation modes.",223            gr.update(value=False, interactive=False),224        )225 226    checkbox_update = gr.update(interactive=True)227    if not enabled:228        return gr.update(visible=False), "Duration control is disabled.", checkbox_update229 230    language, default_tokens, min_tokens, max_tokens = estimate_duration_tokens(text)231    # Slider is initialized with value=1 as a placeholder; treat it as "unset"232    # so first-time estimation uses the computed default instead of clamping to min.233    if current_tokens is None or int(current_tokens) == 1:234        slider_value = default_tokens235    else:236        slider_value = int(current_tokens)237        slider_value = max(min_tokens, min(max_tokens, slider_value))238 239    language_label = "Chinese" if language == "zh" else "English"240    hint = (241        f"Duration control enabled | detected language: {language_label} | "242        f"default={default_tokens}, range=[{min_tokens}, {max_tokens}]"243    )244    return (245        gr.update(246            visible=True,247            minimum=min_tokens,248            maximum=max_tokens,249            value=slider_value,250            step=1,251        ),252        hint,253        checkbox_update,254    )255 256 257def normalize_language_tag(language_tag: str | None) -> str | None:258    language_tag = (language_tag or "").strip()259    if not language_tag or language_tag == LANGUAGE_TAG_AUTO:260        return None261    return language_tag262 263 264def build_conversation(265    text: str,266    reference_audio: str | None,267    mode_with_reference: str,268    expected_tokens: int | None,269    language_tag: str | None,270    processor,271):272    text = (text or "").strip()273    if not text:274        raise ValueError("Please enter text to synthesize.")275 276    user_kwargs = {"text": text}277    normalized_language = normalize_language_tag(language_tag)278    if normalized_language is not None:279        user_kwargs["language"] = normalized_language280    if expected_tokens is not None:281        user_kwargs["tokens"] = int(expected_tokens)282 283    if not reference_audio:284        conversations = [[processor.build_user_message(**user_kwargs)]]285        return conversations, "generation", "Direct Generation"286 287    if mode_with_reference == MODE_CLONE:288        clone_kwargs = dict(user_kwargs)289        clone_kwargs["reference"] = [reference_audio]290        conversations = [[processor.build_user_message(**clone_kwargs)]]291        return conversations, "generation", MODE_CLONE292 293    if mode_with_reference == MODE_CONTINUE:294        conversations = [295            [296                processor.build_user_message(**user_kwargs),297                processor.build_assistant_message(audio_codes_list=[reference_audio]),298            ]299        ]300        return conversations, "continuation", MODE_CONTINUE301 302    continue_clone_kwargs = dict(user_kwargs)303    continue_clone_kwargs["reference"] = [reference_audio]304    conversations = [305        [306            processor.build_user_message(**continue_clone_kwargs),307            processor.build_assistant_message(audio_codes_list=[reference_audio]),308        ]309    ]310    return conversations, "continuation", MODE_CONTINUE_CLONE311 312 313def render_mode_hint(reference_audio: str | None, mode_with_reference: str):314    if not reference_audio:315        return "Current mode: **Direct Generation** (no reference audio uploaded)"316    if mode_with_reference == MODE_CLONE:317        return "Current mode: **Clone** (speaker timbre will be cloned from the reference audio)"318    return f"Current mode: **{mode_with_reference}**  \n> {CONTINUATION_NOTICE}"319 320 321def apply_example_selection(322    mode_with_reference: str,323    duration_control_enabled: bool,324    duration_tokens: int,325    evt: gr.SelectData,326):327    if evt is None or evt.index is None:328        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()329 330    if isinstance(evt.index, (tuple, list)):331        row_idx = int(evt.index[0])332    else:333        row_idx = int(evt.index)334 335    if row_idx < 0 or row_idx >= len(EXAMPLE_ROWS):336        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()337 338    _, audio_path, example_text = EXAMPLE_ROWS[row_idx]339    duration_slider_update, duration_hint, duration_checkbox_update = update_duration_controls(340        duration_control_enabled,341        example_text,342        duration_tokens,343        mode_with_reference,344    )345    return (346        audio_path,347        example_text,348        render_mode_hint(audio_path, mode_with_reference),349        duration_slider_update,350        duration_hint,351        duration_checkbox_update,352    )353 354 355@spaces.GPU(duration=180)356def run_inference(357    text: str,358    reference_audio: str | None,359    mode_with_reference: str,360    duration_control_enabled: bool,361    duration_tokens: int,362    language_tag: str | None,363    temperature: float,364    top_p: float,365    top_k: int,366    repetition_penalty: float,367    model_path: str,368    device: str,369    attn_implementation: str,370    max_new_tokens: int,371):372    started_at = time.monotonic()373    model, processor, torch_device, sample_rate = load_backend(374        model_path=model_path,375        device_str=device,376        attn_implementation=attn_implementation,377    )378    duration_enabled = bool(duration_control_enabled and supports_duration_control(mode_with_reference))379    expected_tokens = int(duration_tokens) if duration_enabled else None380    conversations, mode, mode_name = build_conversation(381        text=text,382        reference_audio=reference_audio,383        mode_with_reference=mode_with_reference,384        expected_tokens=expected_tokens,385        language_tag=language_tag,386        processor=processor,387    )388 389    batch = processor(conversations, mode=mode)390    input_ids = batch["input_ids"].to(torch_device)391    attention_mask = batch["attention_mask"].to(torch_device)392 393    with torch.no_grad():394        outputs = model.generate(395            input_ids=input_ids,396            attention_mask=attention_mask,397            max_new_tokens=int(max_new_tokens),398            audio_temperature=float(temperature),399            audio_top_p=float(top_p),400            audio_top_k=int(top_k),401            audio_repetition_penalty=float(repetition_penalty),402        )403 404    messages = processor.decode(outputs)405    if not messages or messages[0] is None:406        raise RuntimeError("The model did not return a decodable audio result.")407 408    audio = messages[0].audio_codes_list[0]409    if isinstance(audio, torch.Tensor):410        audio_np = audio.detach().float().cpu().numpy()411    else:412        audio_np = np.asarray(audio, dtype=np.float32)413 414    if audio_np.ndim > 1:415        audio_np = audio_np.reshape(-1)416    audio_np = audio_np.astype(np.float32, copy=False)417 418    elapsed = time.monotonic() - started_at419    normalized_language = normalize_language_tag(language_tag)420    status = (421        f"Done | mode: {mode_name} | language={normalized_language or 'auto'} | "422        f"elapsed: {elapsed:.2f}s | "423        f"max_new_tokens={int(max_new_tokens)}, "424        f"expected_tokens={expected_tokens if expected_tokens is not None else 'off'}, "425        f"audio_temperature={float(temperature):.2f}, audio_top_p={float(top_p):.2f}, "426        f"audio_top_k={int(top_k)}, audio_repetition_penalty={float(repetition_penalty):.2f}"427    )428    return (sample_rate, audio_np), status429 430 431def build_demo(args: argparse.Namespace):432    custom_css = """433    :root {434      --bg: #f6f7f8;435      --panel: #ffffff;436      --ink: #111418;437      --muted: #4d5562;438      --line: #e5e7eb;439      --accent: #0f766e;440    }441    .gradio-container {442      background: linear-gradient(180deg, #f7f8fa 0%, #f3f5f7 100%);443      color: var(--ink);444    }445    .app-card {446      border: 1px solid var(--line);447      border-radius: 16px;448      background: var(--panel);449      padding: 14px;450    }451    .app-title {452      font-size: 22px;453      font-weight: 700;454      margin-bottom: 6px;455      letter-spacing: 0.2px;456    }457    .app-subtitle {458      color: var(--muted);459      font-size: 14px;460      margin-bottom: 8px;461    }462    #output_audio {463      padding-bottom: 12px;464      margin-bottom: 8px;465      overflow: hidden !important;466    }467    #output_audio > .wrap {468      overflow: hidden !important;469    }470    #output_audio audio {471      margin-bottom: 6px;472    }473    #run-btn {474      background: var(--accent);475      border: none;476    }477    """478 479    with gr.Blocks(title="MOSS-TTS Demo") as demo:480        gr.Markdown(481            """482            <div class="app-card">483              <div class="app-title">MOSS-TTS v1.5</div>484              <div class="app-subtitle">Direct Generation, Clone, Continuation, Continuation + Clone, language tags, and inline pause markers</div>485            </div>486            """487        )488 489        with gr.Row(equal_height=False):490            with gr.Column(scale=3):491                text = gr.Textbox(492                    label="Text",493                    lines=9,494                    placeholder="Enter text to synthesize. In continuation modes, prepend the reference audio transcript.",495                )496                reference_audio = gr.Audio(497                    label="Reference Audio (Optional)",498                    type="filepath",499                )500                mode_with_reference = gr.Radio(501                    choices=[MODE_CLONE, MODE_CONTINUE, MODE_CONTINUE_CLONE],502                    value=MODE_CLONE,503                    label="Mode with Reference Audio",504                    info="If no reference audio is uploaded, Direct Generation will be used automatically.",505                )506                mode_hint = gr.Markdown(render_mode_hint(None, MODE_CLONE))507                language_tag = gr.Dropdown(508                    choices=LANGUAGE_TAG_CHOICES,509                    value=LANGUAGE_TAG_AUTO,510                    label="Language Tag",511                    info="Optional for v1.5. Set this when the input language is known, especially outside Chinese and English.",512                )513                duration_control_enabled = gr.Checkbox(514                    value=False,515                    label="Enable Duration Control (Expected Audio Tokens)",516                )517                duration_tokens = gr.Slider(518                    minimum=1,519                    maximum=2,520                    step=1,521                    value=1,522                    label="expected_tokens",523                    visible=False,524                )525                duration_hint = gr.Markdown("Duration control is disabled.")526 527                with gr.Accordion("Sampling Parameters (Audio)", open=True):528                    temperature = gr.Slider(529                        minimum=0.1,530                        maximum=3.0,531                        step=0.05,532                        value=1.7,533                        label="temperature",534                    )535                    top_p = gr.Slider(536                        minimum=0.1,537                        maximum=1.0,538                        step=0.01,539                        value=0.8,540                        label="top_p",541                    )542                    top_k = gr.Slider(543                        minimum=1,544                        maximum=200,545                        step=1,546                        value=25,547                        label="top_k",548                    )549                    repetition_penalty = gr.Slider(550                        minimum=0.8,551                        maximum=2.0,552                        step=0.05,553                        value=1.0,554                        label="repetition_penalty",555                    )556                    max_new_tokens = gr.Slider(557                        minimum=256,558                        maximum=8192,559                        step=128,560                        value=DEFAULT_MAX_NEW_TOKENS,561                        label="max_new_tokens",562                    )563 564                run_btn = gr.Button("Generate Speech", variant="primary", elem_id="run-btn")565 566            with gr.Column(scale=2):567                output_audio = gr.Audio(label="Output Audio", type="numpy", elem_id="output_audio")568                status = gr.Textbox(label="Status", lines=4, interactive=False)569                examples_table = gr.Dataframe(570                    headers=["Reference Speech", "Example Text"],571                    value=[[name, text] for name, _, text in EXAMPLE_ROWS],572                    datatype=["str", "str"],573                    row_count=(len(EXAMPLE_ROWS), "fixed"),574                    col_count=(2, "fixed"),575                    interactive=False,576                    wrap=True,577                    label="Examples (click a row to fill inputs)",578                )579 580        reference_audio.change(581            fn=render_mode_hint,582            inputs=[reference_audio, mode_with_reference],583            outputs=[mode_hint],584        )585        mode_with_reference.change(586            fn=render_mode_hint,587            inputs=[reference_audio, mode_with_reference],588            outputs=[mode_hint],589        )590        duration_control_enabled.change(591            fn=update_duration_controls,592            inputs=[duration_control_enabled, text, duration_tokens, mode_with_reference],593            outputs=[duration_tokens, duration_hint, duration_control_enabled],594        )595        text.change(596            fn=update_duration_controls,597            inputs=[duration_control_enabled, text, duration_tokens, mode_with_reference],598            outputs=[duration_tokens, duration_hint, duration_control_enabled],599        )600        mode_with_reference.change(601            fn=update_duration_controls,602            inputs=[duration_control_enabled, text, duration_tokens, mode_with_reference],603            outputs=[duration_tokens, duration_hint, duration_control_enabled],604        )605        examples_table.select(606            fn=apply_example_selection,607            inputs=[mode_with_reference, duration_control_enabled, duration_tokens],608            outputs=[609                reference_audio,610                text,611                mode_hint,612                duration_tokens,613                duration_hint,614                duration_control_enabled,615            ],616        )617 618        run_btn.click(619            fn=run_inference,620            inputs=[621                text,622                reference_audio,623                mode_with_reference,624                duration_control_enabled,625                duration_tokens,626                language_tag,627                temperature,628                top_p,629                top_k,630                repetition_penalty,631                gr.State(args.model_path),632                gr.State(args.device),633                gr.State(args.attn_implementation),634                max_new_tokens,635            ],636            outputs=[output_audio, status],637        )638    demo._moss_custom_css = custom_css639    return demo640 641 642def resolve_runtime_attn(args: argparse.Namespace) -> argparse.Namespace:643    runtime_device = torch.device(args.device if torch.cuda.is_available() else "cpu")644    runtime_dtype = torch.bfloat16 if runtime_device.type == "cuda" else torch.float32645    args.attn_implementation = resolve_attn_implementation(646        requested=args.attn_implementation,647        device=runtime_device,648        dtype=runtime_dtype,649    ) or "none"650    return args651 652 653def parse_bool_env(name: str, default: bool) -> bool:654    value = os.getenv(name)655    if value is None:656        return default657    return value.strip().lower() in {"1", "true", "yes", "y", "on"}658 659 660def parse_port(value: str | None, default: int) -> int:661    if not value:662        return default663    try:664        return int(value)665    except ValueError:666        return default667 668 669def main():670    parser = argparse.ArgumentParser(description="MossTTS Gradio Demo")671    parser.add_argument("--model_path", type=str, default=MODEL_PATH)672    parser.add_argument("--device", type=str, default="cuda:0")673    parser.add_argument("--attn_implementation", type=str, default=DEFAULT_ATTN_IMPLEMENTATION)674    parser.add_argument("--host", type=str, default="0.0.0.0")675    parser.add_argument(676        "--port",677        type=int,678        default=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),679    )680    parser.add_argument("--share", action="store_true")681    args = parser.parse_args()682 683    args.host = os.getenv("GRADIO_SERVER_NAME", args.host)684    args.port = parse_port(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT")), args.port)685    args = resolve_runtime_attn(args)686    print(f"[INFO] Using attn_implementation={args.attn_implementation}", flush=True)687 688    preload_enabled = parse_bool_env(PRELOAD_ENV_VAR, default=not bool(os.getenv("SPACE_ID")))689    if preload_enabled:690        preload_started_at = time.monotonic()691        print(692            f"[Startup] Preloading backend: model={args.model_path}, device={args.device}, attn={args.attn_implementation}",693            flush=True,694        )695        load_backend(696            model_path=args.model_path,697            device_str=args.device,698            attn_implementation=args.attn_implementation,699        )700        print(701            f"[Startup] Backend preload finished in {time.monotonic() - preload_started_at:.2f}s",702            flush=True,703        )704    else:705        print(706            f"[Startup] Skipping preload (set {PRELOAD_ENV_VAR}=1 to enable).",707            flush=True,708        )709 710    demo = build_demo(args)711    launch_kwargs = {712        "server_name": args.host,713        "server_port": args.port,714        "share": args.share,715    }716    launch_parameters = inspect.signature(demo.launch).parameters717    if "css" in launch_parameters:718        launch_kwargs["css"] = getattr(demo, "_moss_custom_css", None)719    if "ssr_mode" in launch_parameters:720        launch_kwargs["ssr_mode"] = False721 722    demo.queue(max_size=16, default_concurrency_limit=1).launch(**launch_kwargs)723 724 725if __name__ == "__main__":726    main()727