| from __future__ import annotations |
|
|
| import contextlib |
| import io |
| import os |
| from pathlib import Path |
| from typing import List, Tuple |
|
|
| import gradio as gr |
| import torch |
| import spaces |
|
|
| from dia2 import Dia2, GenerationConfig, SamplingConfig |
|
|
| DEFAULT_REPO = os.environ.get("DIA2_DEFAULT_REPO", "nari-labs/Dia2-2B") |
| MAX_TURNS = 10 |
| INITIAL_TURNS = 2 |
|
|
| _dia: Dia2 | None = None |
|
|
|
|
| def _get_dia() -> Dia2: |
| global _dia |
| if _dia is None: |
| _dia = Dia2.from_repo(DEFAULT_REPO, device="cuda", dtype="bfloat16") |
| return _dia |
|
|
|
|
| def _concat_script(turn_count: int, turn_values: List[str]) -> str: |
| lines: List[str] = [] |
| for idx in range(min(turn_count, len(turn_values))): |
| text = (turn_values[idx] or "").strip() |
| if not text: |
| continue |
| speaker = "[S1]" if idx % 2 == 0 else "[S2]" |
| lines.append(f"{speaker} {text}") |
| return "\n".join(lines) |
|
|
|
|
| EXAMPLES: dict[str, dict[str, List[str] | str | None]] = { |
| "Intro": { |
| "turns": [ |
| "Hello Dia2 fans! Today we're unveiling the new open TTS model.", |
| "Sounds exciting. Can you show a sample right now?", |
| "Absolutely. (laughs) Just press generate.", |
| ], |
| "voice_s1": "example_prefix1.wav", |
| "voice_s2": "example_prefix2.wav", |
| }, |
| "Customer Support": { |
| "turns": [ |
| "Thanks for calling. How can I help you today?", |
| "My parcel never arrived and it's been two weeks.", |
| "I'm sorry about that. Let me check your tracking number.", |
| "Appreciate it. I really need that package soon.", |
| ], |
| "voice_s1": "example_prefix1.wav", |
| "voice_s2": "example_prefix2.wav", |
| }, |
| } |
|
|
|
|
| def _apply_turn_visibility(count: int) -> List[gr.Update]: |
| return [gr.update(visible=i < count) for i in range(MAX_TURNS)] |
|
|
|
|
| def _add_turn(count: int): |
| count = min(count + 1, MAX_TURNS) |
| return (count, *_apply_turn_visibility(count)) |
|
|
|
|
| def _remove_turn(count: int): |
| count = max(1, count - 1) |
| return (count, *_apply_turn_visibility(count)) |
|
|
|
|
| def _load_example(name: str, count: int): |
| data = EXAMPLES.get(name) |
| if not data: |
| return (count, *_apply_turn_visibility(count), None, None) |
| turns = data.get("turns", []) |
| voice_s1_path = data.get("voice_s1") |
| voice_s2_path = data.get("voice_s2") |
| new_count = min(len(turns), MAX_TURNS) |
| updates: List[gr.Update] = [] |
| for idx in range(MAX_TURNS): |
| if idx < new_count: |
| updates.append(gr.update(value=turns[idx], visible=True)) |
| else: |
| updates.append(gr.update(value="", visible=idx < INITIAL_TURNS)) |
| return (new_count, *updates, voice_s1_path, voice_s2_path) |
|
|
|
|
| def _prepare_prefix(file_path: str | None) -> str | None: |
| if not file_path: |
| return None |
| path = Path(file_path) |
| if not path.exists(): |
| return None |
| return str(path) |
|
|
|
|
| @spaces.GPU(duration=100) |
| def generate_audio( |
| turn_count: int, |
| *inputs, |
| ): |
| turn_values = list(inputs[:MAX_TURNS]) |
| voice_s1 = inputs[MAX_TURNS] |
| voice_s2 = inputs[MAX_TURNS + 1] |
| cfg_scale = float(inputs[MAX_TURNS + 2]) |
| text_temperature = float(inputs[MAX_TURNS + 3]) |
| audio_temperature = float(inputs[MAX_TURNS + 4]) |
| text_top_k = int(inputs[MAX_TURNS + 5]) |
| audio_top_k = int(inputs[MAX_TURNS + 6]) |
| include_prefix = bool(inputs[MAX_TURNS + 7]) |
|
|
| script = _concat_script(turn_count, turn_values) |
| if not script.strip(): |
| raise gr.Error("Please enter at least one non-empty speaker turn.") |
|
|
| dia = _get_dia() |
| config = GenerationConfig( |
| cfg_scale=cfg_scale, |
| text=SamplingConfig(temperature=text_temperature, top_k=text_top_k), |
| audio=SamplingConfig(temperature=audio_temperature, top_k=audio_top_k), |
| use_cuda_graph=True, |
| ) |
| kwargs = { |
| "prefix_speaker_1": _prepare_prefix(voice_s1), |
| "prefix_speaker_2": _prepare_prefix(voice_s2), |
| "include_prefix": include_prefix, |
| } |
| buffer = io.StringIO() |
| with contextlib.redirect_stdout(buffer): |
| result = dia.generate( |
| script, |
| config=config, |
| output_wav=None, |
| verbose=True, |
| **kwargs, |
| ) |
| waveform = result.waveform.detach().cpu().numpy() |
| sample_rate = result.sample_rate |
| timestamps = result.timestamps |
| log_text = buffer.getvalue().strip() |
| table = [[w, round(t, 3)] for w, t in timestamps] |
| return (sample_rate, waveform), table, log_text or "Generation finished." |
|
|
|
|
| def build_interface() -> gr.Blocks: |
| with gr.Blocks( |
| title="Dia2 TTS", css=".compact-turn textarea {min-height: 60px}" |
| ) as demo: |
| gr.Markdown( |
| """## Dia2 — Open TTS Model |
| Compose dialogue, attach optional voice prompts, and generate audio (CUDA graphs enabled by default).""" |
| ) |
| turn_state = gr.State(INITIAL_TURNS) |
| with gr.Row(equal_height=True): |
| example_dropdown = gr.Dropdown( |
| choices=["(select example)"] + list(EXAMPLES.keys()), |
| label="Examples", |
| value="(select example)", |
| ) |
| with gr.Row(equal_height=True): |
| with gr.Column(scale=1): |
| with gr.Group(): |
| gr.Markdown("### Script") |
| controls = [] |
| for idx in range(MAX_TURNS): |
| speaker = "[S1]" if idx % 2 == 0 else "[S2]" |
| box = gr.Textbox( |
| label=f"{speaker} turn {idx + 1}", |
| lines=2, |
| elem_classes=["compact-turn"], |
| placeholder=f"Enter dialogue for {speaker}…", |
| visible=idx < INITIAL_TURNS, |
| ) |
| controls.append(box) |
| with gr.Row(): |
| add_btn = gr.Button("Add Turn") |
| remove_btn = gr.Button("Remove Turn") |
| with gr.Group(): |
| gr.Markdown("### Voice Prompts") |
| with gr.Row(): |
| voice_s1 = gr.File( |
| label="[S1] voice (wav/mp3)", type="filepath" |
| ) |
| voice_s2 = gr.File( |
| label="[S2] voice (wav/mp3)", type="filepath" |
| ) |
| with gr.Group(): |
| gr.Markdown("### Sampling") |
| cfg_scale = gr.Slider( |
| 1.0, 8.0, value=6.0, step=0.1, label="CFG Scale" |
| ) |
| with gr.Group(): |
| gr.Markdown("#### Text Sampling") |
| text_temperature = gr.Slider( |
| 0.1, 1.5, value=0.6, step=0.05, label="Text Temperature" |
| ) |
| text_top_k = gr.Slider( |
| 1, 200, value=50, step=1, label="Text Top-K" |
| ) |
| with gr.Group(): |
| gr.Markdown("#### Audio Sampling") |
| audio_temperature = gr.Slider( |
| 0.1, 1.5, value=0.8, step=0.05, label="Audio Temperature" |
| ) |
| audio_top_k = gr.Slider( |
| 1, 200, value=50, step=1, label="Audio Top-K" |
| ) |
| include_prefix = gr.Checkbox( |
| label="Keep prefix audio in output", value=False |
| ) |
| generate_btn = gr.Button("Generate", variant="primary") |
| with gr.Column(scale=1): |
| gr.Markdown("### Output") |
| audio_out = gr.Audio(label="Waveform", interactive=False) |
| timestamps = gr.Dataframe( |
| headers=["word", "seconds"], label="Timestamps" |
| ) |
| log_box = gr.Textbox(label="Logs", lines=8) |
|
|
| add_btn.click( |
| lambda c: _add_turn(c), |
| inputs=turn_state, |
| outputs=[turn_state, *controls], |
| ) |
| remove_btn.click( |
| lambda c: _remove_turn(c), |
| inputs=turn_state, |
| outputs=[turn_state, *controls], |
| ) |
| example_dropdown.change( |
| lambda name, c: _load_example(name, c), |
| inputs=[example_dropdown, turn_state], |
| outputs=[turn_state, *controls, voice_s1, voice_s2], |
| ) |
|
|
| generate_btn.click( |
| generate_audio, |
| inputs=[ |
| turn_state, |
| *controls, |
| voice_s1, |
| voice_s2, |
| cfg_scale, |
| text_temperature, |
| audio_temperature, |
| text_top_k, |
| audio_top_k, |
| include_prefix, |
| ], |
| outputs=[audio_out, timestamps, log_box], |
| ) |
| return demo |
|
|
|
|
| if __name__ == "__main__": |
| app = build_interface() |
| app.queue(default_concurrency_limit=1) |
| app.launch(share=True) |
|
|
