Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,264 @@
+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)