OpenMOSS-Team/MOSS-VoiceGenerator

1import argparse2import functools3import importlib.util4import json5import os6from pathlib import Path7import re8import time9 10try:11    import spaces12except ImportError:13    class _SpacesFallback:14        @staticmethod15        def GPU(*_args, **_kwargs):16            def _decorator(func):17                return func18 19            return _decorator20 21    spaces = _SpacesFallback()22 23import gradio as gr24import numpy as np25import torch26from transformers import AutoModel, AutoProcessor27 28# Disable the broken cuDNN SDPA backend29torch.backends.cuda.enable_cudnn_sdp(False)30# Keep these enabled as fallbacks31torch.backends.cuda.enable_flash_sdp(True)32torch.backends.cuda.enable_mem_efficient_sdp(True)33torch.backends.cuda.enable_math_sdp(True)34 35MODEL_PATH = "OpenMOSS-Team/MOSS-VoiceGenerator"36DEFAULT_ATTN_IMPLEMENTATION = "auto"37DEFAULT_MAX_NEW_TOKENS = 409638PRELOAD_ENV_VAR = "MOSS_VOICE_GENERATOR_PRELOAD_AT_STARTUP"39EXAMPLE_TEXTS_JSONL_PATH = Path(__file__).resolve().parent / "text" / "moss_voice_generator_example_texts.jsonl"40 41 42def _parse_example_id(example_id: str) -> tuple[str, int] | None:43    matched = re.fullmatch(r"(zh|en)/(\d+)", (example_id or "").strip())44    if matched is None:45        return None46    return matched.group(1), int(matched.group(2))47 48 49def build_example_rows() -> list[tuple[str, str, str]]:50    rows: list[tuple[str, int, str, str]] = []51    with open(EXAMPLE_TEXTS_JSONL_PATH, "r", encoding="utf-8") as f:52        for line in f:53            if not line.strip():54                continue55            sample = json.loads(line)56            parsed = _parse_example_id(sample.get("id", ""))57            if parsed is None:58                continue59 60            language, index = parsed61            instruction = str(sample.get("instruction", "")).strip()62            text = str(sample.get("text", "")).strip()63            rows.append((language, index, instruction, text))64 65    language_order = {"zh": 0, "en": 1}66    rows.sort(key=lambda item: (language_order.get(item[0], 99), item[1]))67    return [(f"{language}/{index}", instruction, text) for language, index, instruction, text in rows]68 69 70EXAMPLE_ROWS = build_example_rows()71 72 73def apply_example_selection(evt: gr.SelectData):74    if evt is None or evt.index is None:75        return gr.update(), gr.update()76 77    if isinstance(evt.index, (tuple, list)):78        row_idx = int(evt.index[0])79    else:80        row_idx = int(evt.index)81 82    if row_idx < 0 or row_idx >= len(EXAMPLE_ROWS):83        return gr.update(), gr.update()84 85    _, instruction_value, text_value = EXAMPLE_ROWS[row_idx]86    return instruction_value, text_value87 88 89def resolve_attn_implementation(requested: str, device: torch.device, dtype: torch.dtype) -> str | None:90    requested_norm = (requested or "").strip().lower()91 92    if requested_norm in {"none"}:93        return None94 95    if requested_norm not in {"", "auto"}:96        return requested97 98    # Prefer FlashAttention 2 when package + device conditions are met.99    if (100        device.type == "cuda"101        and importlib.util.find_spec("flash_attn") is not None102        and dtype in {torch.float16, torch.bfloat16}103    ):104        major, _ = torch.cuda.get_device_capability(device)105        if major >= 8:106            return "flash_attention_2"107 108    # CUDA fallback: use PyTorch SDPA kernels.109    if device.type == "cuda":110        return "sdpa"111 112    # CPU fallback.113    return "eager"114 115 116@functools.lru_cache(maxsize=1)117def load_backend(model_path: str, device_str: str, attn_implementation: str):118    device = torch.device(device_str if torch.cuda.is_available() else "cpu")119    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32120    resolved_attn_implementation = resolve_attn_implementation(121        requested=attn_implementation,122        device=device,123        dtype=dtype,124    )125 126    processor = AutoProcessor.from_pretrained(127        model_path,128        trust_remote_code=True,129        normalize_inputs=True,130    )131    if hasattr(processor, "audio_tokenizer"):132        processor.audio_tokenizer = processor.audio_tokenizer.to(device)133        processor.audio_tokenizer.eval()134 135    model_kwargs = {136        "trust_remote_code": True,137        "torch_dtype": dtype,138    }139    if resolved_attn_implementation:140        model_kwargs["attn_implementation"] = resolved_attn_implementation141 142    model = AutoModel.from_pretrained(model_path, **model_kwargs).to(device)143    model.eval()144 145    sample_rate = int(getattr(processor.model_config, "sampling_rate", 24000))146    return model, processor, device, sample_rate147 148 149def build_conversation(text: str, instruction: str, processor):150    text = (text or "").strip()151    instruction = (instruction or "").strip()152    if not text:153        raise ValueError("Please enter text to synthesize.")154    if not instruction:155        raise ValueError("Please enter a voice instruction.")156 157    return [[processor.build_user_message(text=text, instruction=instruction)]]158 159 160@spaces.GPU(duration=180)161def run_inference(162    text: str,163    instruction: str,164    temperature: float,165    top_p: float,166    top_k: int,167    repetition_penalty: float,168    max_new_tokens: int,169    model_path: str,170    device: str,171    attn_implementation: str,172):173    started_at = time.monotonic()174    model, processor, torch_device, sample_rate = load_backend(175        model_path=model_path,176        device_str=device,177        attn_implementation=attn_implementation,178    )179 180    conversations = build_conversation(181        text=text,182        instruction=instruction,183        processor=processor,184    )185 186    batch = processor(conversations, mode="generation")187    input_ids = batch["input_ids"].to(torch_device)188    attention_mask = batch["attention_mask"].to(torch_device)189 190    with torch.no_grad():191        outputs = model.generate(192            input_ids=input_ids,193            attention_mask=attention_mask,194            max_new_tokens=int(max_new_tokens),195            audio_temperature=float(temperature),196            audio_top_p=float(top_p),197            audio_top_k=int(top_k),198            audio_repetition_penalty=float(repetition_penalty),199        )200 201    messages = processor.decode(outputs)202    if not messages or messages[0] is None:203        raise RuntimeError("The model did not return a decodable audio result.")204 205    audio = messages[0].audio_codes_list[0]206    if isinstance(audio, torch.Tensor):207        audio_np = audio.detach().float().cpu().numpy()208    else:209        audio_np = np.asarray(audio, dtype=np.float32)210 211    if audio_np.ndim > 1:212        audio_np = audio_np.reshape(-1)213    audio_np = audio_np.astype(np.float32, copy=False)214 215    elapsed = time.monotonic() - started_at216    status = (217        f"Done | elapsed: {elapsed:.2f}s | "218        f"max_new_tokens={int(max_new_tokens)}, "219        f"audio_temperature={float(temperature):.2f}, audio_top_p={float(top_p):.2f}, "220        f"audio_top_k={int(top_k)}, audio_repetition_penalty={float(repetition_penalty):.2f}"221    )222    return (sample_rate, audio_np), status223 224 225def build_demo(args: argparse.Namespace):226    custom_css = """227    :root {228      --bg: #f6f7f8;229      --panel: #ffffff;230      --ink: #111418;231      --muted: #4d5562;232      --line: #e5e7eb;233      --accent: #0f766e;234    }235    .gradio-container {236      background: linear-gradient(180deg, #f7f8fa 0%, #f3f5f7 100%);237      color: var(--ink);238    }239    .app-card {240      border: 1px solid var(--line);241      border-radius: 16px;242      background: var(--panel);243      padding: 14px;244    }245    .app-title {246      font-size: 22px;247      font-weight: 700;248      margin-bottom: 6px;249      letter-spacing: 0.2px;250    }251    .app-subtitle {252      color: var(--muted);253      font-size: 14px;254      margin-bottom: 8px;255    }256    #output_audio {257      padding-bottom: 12px;258      margin-bottom: 8px;259      overflow: hidden !important;260    }261    #output_audio > .wrap {262      overflow: hidden !important;263    }264    #output_audio audio {265      margin-bottom: 6px;266    }267    #run-btn {268      background: var(--accent);269      border: none;270    }271    """272 273    with gr.Blocks(title="MOSS-VoiceGenerator Demo", css=custom_css) as demo:274        gr.Markdown(275            """276            <div class="app-card">277              <div class="app-title">MOSS-VoiceGenerator</div>278              <div class="app-subtitle">Design expressive voices from instruction + text without reference audio.</div>279            </div>280            """281        )282 283        with gr.Row(equal_height=False):284            with gr.Column(scale=3):285                instruction = gr.Textbox(286                    label="Voice Instruction",287                    lines=5,288                    placeholder="Example: Warm, gentle female narrator voice with calm pacing and clear articulation.",289                )290                text = gr.Textbox(291                    label="Text",292                    lines=8,293                    placeholder="Enter the text content to synthesize with the instruction-defined voice.",294                )295 296                with gr.Accordion("Sampling Parameters (Audio)", open=True):297                    temperature = gr.Slider(298                        minimum=0.1,299                        maximum=3.0,300                        step=0.05,301                        value=1.5,302                        label="temperature",303                    )304                    top_p = gr.Slider(305                        minimum=0.1,306                        maximum=1.0,307                        step=0.01,308                        value=0.6,309                        label="top_p",310                    )311                    top_k = gr.Slider(312                        minimum=1,313                        maximum=200,314                        step=1,315                        value=50,316                        label="top_k",317                    )318                    repetition_penalty = gr.Slider(319                        minimum=0.8,320                        maximum=2.0,321                        step=0.05,322                        value=1.1,323                        label="repetition_penalty",324                    )325                    max_new_tokens = gr.Slider(326                        minimum=256,327                        maximum=8192,328                        step=128,329                        value=DEFAULT_MAX_NEW_TOKENS,330                        label="max_new_tokens",331                    )332 333                run_btn = gr.Button("Generate Voice", variant="primary", elem_id="run-btn")334 335            with gr.Column(scale=2):336                output_audio = gr.Audio(label="Output Audio", type="numpy", elem_id="output_audio")337                status = gr.Textbox(label="Status", lines=4, interactive=False)338                examples_table = gr.Dataframe(339                    headers=["Voice Instruction", "Example Text"],340                    value=[[example_instruction, example_text] for _, example_instruction, example_text in EXAMPLE_ROWS],341                    datatype=["str", "str"],342                    row_count=(len(EXAMPLE_ROWS), "fixed"),343                    col_count=(2, "fixed"),344                    interactive=False,345                    wrap=True,346                    label="Examples (click a row to fill inputs)",347                )348 349        examples_table.select(350            fn=apply_example_selection,351            inputs=[],352            outputs=[instruction, text],353        )354 355        run_btn.click(356            fn=run_inference,357            inputs=[358                text,359                instruction,360                temperature,361                top_p,362                top_k,363                repetition_penalty,364                max_new_tokens,365                gr.State(args.model_path),366                gr.State(args.device),367                gr.State(args.attn_implementation),368            ],369            outputs=[output_audio, status],370        )371    return demo372 373 374def resolve_runtime_attn(args: argparse.Namespace) -> argparse.Namespace:375    runtime_device = torch.device(args.device if torch.cuda.is_available() else "cpu")376    runtime_dtype = torch.bfloat16 if runtime_device.type == "cuda" else torch.float32377    args.attn_implementation = resolve_attn_implementation(378        requested=args.attn_implementation,379        device=runtime_device,380        dtype=runtime_dtype,381    ) or "none"382    return args383 384 385def parse_bool_env(name: str, default: bool) -> bool:386    value = os.getenv(name)387    if value is None:388        return default389    return value.strip().lower() in {"1", "true", "yes", "y", "on"}390 391 392def parse_port(value: str | None, default: int) -> int:393    if not value:394        return default395    try:396        return int(value)397    except ValueError:398        return default399 400 401def main():402    parser = argparse.ArgumentParser(description="MOSS-VoiceGenerator Gradio Demo")403    parser.add_argument("--model_path", type=str, default=MODEL_PATH)404    parser.add_argument("--device", type=str, default="cuda:0")405    parser.add_argument("--attn_implementation", type=str, default=DEFAULT_ATTN_IMPLEMENTATION)406    parser.add_argument("--host", type=str, default="0.0.0.0")407    parser.add_argument(408        "--port",409        type=int,410        default=int(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT", "7860"))),411    )412    parser.add_argument("--share", action="store_true")413    args = parser.parse_args()414 415    args.host = os.getenv("GRADIO_SERVER_NAME", args.host)416    args.port = parse_port(os.getenv("GRADIO_SERVER_PORT", os.getenv("PORT")), args.port)417    args = resolve_runtime_attn(args)418    print(f"[INFO] Using attn_implementation={args.attn_implementation}", flush=True)419 420    preload_enabled = parse_bool_env(PRELOAD_ENV_VAR, default=not bool(os.getenv("SPACE_ID")))421    if preload_enabled:422        preload_started_at = time.monotonic()423        print(424            f"[Startup] Preloading backend: model={args.model_path}, device={args.device}, attn={args.attn_implementation}",425            flush=True,426        )427        load_backend(428            model_path=args.model_path,429            device_str=args.device,430            attn_implementation=args.attn_implementation,431        )432        print(433            f"[Startup] Backend preload finished in {time.monotonic() - preload_started_at:.2f}s",434            flush=True,435        )436    else:437        print(438            f"[Startup] Skipping preload (set {PRELOAD_ENV_VAR}=1 to enable).",439            flush=True,440        )441 442    demo = build_demo(args)443    demo.queue(max_size=16, default_concurrency_limit=1).launch(444        server_name=args.host,445        server_port=args.port,446        share=args.share,447        ssr_mode=False,448    )449 450 451if __name__ == "__main__":452    main()453