app.py

import argparse
import json
import os
import uuid
from pathlib import Path

# Disable PyTorch dynamo/inductor globally for HuggingFace ZeroGPU.
os.environ["TORCHDYNAMO_DISABLE"] = "1"
os.environ["TORCHINDUCTOR_DISABLE"] = "1"

import torch._dynamo as dynamo

dynamo.config.suppress_errors = True

import gradio as gr
import numpy as np
import soundfile as sf
import spaces
import torch

from voxtream.config import SpeechGeneratorConfig
from voxtream.generator import SpeechGenerator
from voxtream.utils.app import (
    CUSTOM_CSS,
    AppConfig,
    GenerationControl,
    SharedGenerationState,
    SpeakingRateState,
    VisualizationState,
    build_low_latency_audio_head,
    clear_outputs,
    empty_rate_plot,
    float32_to_int16,
    load_app_config,
    render_audio_stream,
    render_text_progress,
)
from voxtream.utils.generator import DTYPE_MAP, existing_file, text_generator
from voxtream.utils.generator.helpers import autocast_ctx
from voxtream.utils.generator.text import build_text_progress_metadata


def generation_button_updates(running: bool, paused: bool = False):
    if not running:
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
        )
    if paused:
        return (
            gr.update(visible=False),
            gr.update(visible=True, interactive=True),
            gr.update(visible=True, interactive=True),
        )
    return (
        gr.update(visible=True, interactive=True),
        gr.update(visible=False),
        gr.update(visible=False),
    )


def ensure_generator_on_cuda(speech_generator: SpeechGenerator) -> None:
    if not torch.cuda.is_available():
        return
    if next(speech_generator.model.parameters()).device.type != "cpu":
        return

    dtype = DTYPE_MAP["cuda"]
    speech_generator.model.to("cuda", dtype=dtype)
    speech_generator.mimi.to("cuda", dtype=dtype)
    speech_generator.ctx.mimi_prompt.to("cuda", dtype=dtype)
    speech_generator.ctx.spk_enc.to("cuda", dtype=dtype)
    speech_generator.ctx.device = "cuda"
    speech_generator.ctx.dtype = dtype
    speech_generator._autocast_ctx = autocast_ctx(device="cuda", dtype=dtype)

    if speech_generator._mimi_streaming_started:
        speech_generator._mimi_stream_ctx.__exit__(None, None, None)
        speech_generator._mimi_stream_ctx = None
        speech_generator._mimi_streaming_started = False


def demo_app(
    config: SpeechGeneratorConfig,
    app_config: AppConfig,
    demo_examples,
    synthesize_fn,
    speaking_rate_state: SpeakingRateState,
    generation_control: GenerationControl,
    shared_generation_state: SharedGenerationState,
):
    with gr.Blocks(
        css=CUSTOM_CSS,
        head=build_low_latency_audio_head(app_config),
        title="VoXtream2",
    ) as demo:
        gr.Markdown("# VoXtream2 TTS demo")
        gr.Markdown(
            "⚠️ The initial latency can be high due to deployment on ZeroGPU. "
            "For faster inference, please try local deployment."
        )

        with gr.Row(equal_height=True, elem_id="cols"):
            with gr.Column(scale=1, elem_id="left-col"):
                prompt_audio = gr.Audio(
                    sources=["microphone", "upload"],
                    type="filepath",
                    label=(
                        "Prompt audio (3-10 sec of target voice. "
                        f"Max {config.max_prompt_sec} sec)"
                    ),
                )
                with gr.Accordion("Advanced options", open=False):
                    enable_speaking_rate = gr.Checkbox(
                        label="Use speaking rate control", value=True
                    )
                    prompt_enhancement = gr.Checkbox(
                        label="Prompt enhancement", value=False
                    )
                    prompt_enhancement_msg = gr.Markdown(
                        "⚠️ First 3-5 runs may have higher latency due to model "
                        "loading and warmup.",
                        visible=False,
                    )
                    voice_activity_detection = gr.Checkbox(
                        label="Voice activity detection", value=False
                    )
                    streaming_input = gr.Checkbox(label="Streaming input", value=False)

            with gr.Column(scale=1, elem_id="right-col"):
                target_text = gr.Textbox(
                    lines=4,
                    max_length=config.max_phone_tokens,
                    label=(
                        "Target text (Required, "
                        f"max {config.max_phone_tokens} chars)"
                    ),
                    placeholder="What you want the model to say",
                )
                output_audio = gr.Audio(
                    label="Synthesized audio",
                    interactive=False,
                    streaming=False,
                    autoplay=False,
                    show_download_button=True,
                    show_share_button=False,
                    visible=False,
                )
                stream_audio = gr.HTML(
                    render_audio_stream(app_config), elem_id="audio-stream-container"
                )

        with gr.Row():
            clear_btn = gr.Button("Clear", elem_id="clear", variant="secondary")
            submit_btn = gr.Button(
                "Submit", elem_id="submit", variant="primary", interactive=False
            )
            pause_btn = gr.Button(
                "Pause", elem_id="pause", variant="secondary", visible=False
            )
            resume_btn = gr.Button(
                "Resume", elem_id="resume", variant="primary", visible=False
            )
            stop_btn = gr.Button("Stop", elem_id="stop", variant="stop", visible=False)

        validation_msg = gr.Markdown("", visible=False)
        speaking_rate_control = gr.Slider(
            minimum=app_config.speaking_rate_min,
            maximum=app_config.speaking_rate_max,
            step=app_config.speaking_rate_step,
            value=app_config.speaking_rate_default,
            label="Speaking rate (SPS). Change the speed of speech synthesis in real-time. ",
        )
        rate_plot = gr.HTML(empty_rate_plot(app_config), elem_id="rate-plot-container")
        text_progress = gr.HTML(
            render_text_progress(app_config, None), elem_id="text-progress-container"
        )
        generation_session = gr.State("")

        def validate_inputs(audio, ttext):
            if not audio:
                return gr.update(
                    visible=True, value="⚠️ Please provide a prompt audio."
                ), gr.update(interactive=False)
            if not ttext or not ttext.strip():
                return gr.update(
                    visible=True, value="⚠️ Please provide target text."
                ), gr.update(interactive=False)
            return gr.update(visible=False, value=""), gr.update(interactive=True)

        enable_speaking_rate.change(
            fn=lambda enabled: gr.update(interactive=enabled),
            inputs=enable_speaking_rate,
            outputs=speaking_rate_control,
        )
        prompt_enhancement.change(
            fn=lambda enabled: gr.update(visible=enabled),
            inputs=prompt_enhancement,
            outputs=prompt_enhancement_msg,
        )

        def update_speaking_rate(value, session_id):
            speaking_rate_state.update(value)
            shared_generation_state.update_speaking_rate(session_id, value)

        speaking_rate_control.input(
            fn=update_speaking_rate,
            inputs=[speaking_rate_control, generation_session],
            queue=False,
            show_progress="hidden",
        )
        speaking_rate_control.release(
            fn=update_speaking_rate,
            inputs=[speaking_rate_control, generation_session],
            queue=False,
            show_progress="hidden",
        )

        for inp in [prompt_audio, target_text]:
            inp.change(
                fn=validate_inputs,
                inputs=[prompt_audio, target_text],
                outputs=[validation_msg, submit_btn],
            )

        def prepare_generation(speaking_rate, enable_rate):
            session_id = shared_generation_state.create(speaking_rate)
            generation_control.start()
            speaking_rate_state.start(speaking_rate)
            return (
                gr.update(value=None, visible=False),
                gr.update(interactive=False),
                empty_rate_plot(app_config, show_target=enable_rate),
                render_text_progress(app_config, None),
                render_audio_stream(app_config, session_id=session_id),
                *generation_button_updates(running=True),
                session_id,
            )

        submit_btn.click(
            fn=prepare_generation,
            inputs=[speaking_rate_control, enable_speaking_rate],
            outputs=[
                output_audio,
                enable_speaking_rate,
                rate_plot,
                text_progress,
                stream_audio,
                pause_btn,
                resume_btn,
                stop_btn,
                generation_session,
            ],
            show_progress="hidden",
        ).then(
            fn=synthesize_fn,
            inputs=[
                prompt_audio,
                target_text,
                prompt_enhancement,
                voice_activity_detection,
                streaming_input,
                speaking_rate_control,
                enable_speaking_rate,
                generation_session,
            ],
            outputs=[
                output_audio,
                enable_speaking_rate,
                rate_plot,
                text_progress,
                stream_audio,
                pause_btn,
                resume_btn,
                stop_btn,
                generation_session,
            ],
        )

        def pause_generation(session_id):
            generation_control.pause()
            shared_generation_state.pause(session_id)
            return generation_button_updates(running=True, paused=True)

        def resume_generation(session_id):
            generation_control.resume()
            shared_generation_state.resume(session_id)
            return generation_button_updates(running=True)

        def stop_generation(session_id):
            generation_control.stop()
            speaking_rate_state.stop()
            shared_generation_state.stop(session_id)
            return generation_button_updates(running=False)

        pause_btn.click(
            fn=pause_generation,
            inputs=generation_session,
            outputs=[pause_btn, resume_btn, stop_btn],
            js=(
                "() => { if (window.voxtreamLowLatencyAudio) { "
                "window.voxtreamLowLatencyAudio.pause(); } return []; }"
            ),
            queue=False,
        )
        resume_btn.click(
            fn=resume_generation,
            inputs=generation_session,
            outputs=[pause_btn, resume_btn, stop_btn],
            js=(
                "() => { if (window.voxtreamLowLatencyAudio) { "
                "window.voxtreamLowLatencyAudio.resume(); } return []; }"
            ),
            queue=False,
        )
        stop_btn.click(
            fn=stop_generation,
            inputs=generation_session,
            outputs=[pause_btn, resume_btn, stop_btn],
            js=(
                "() => { if (window.voxtreamLowLatencyAudio) { "
                "window.voxtreamLowLatencyAudio.stop(); } return []; }"
            ),
            queue=False,
        )

        def clear_generation(session_id):
            generation_control.stop()
            speaking_rate_state.stop()
            shared_generation_state.stop(session_id)
            return (
                gr.update(value=None),
                gr.update(value=""),
                gr.update(value=None, visible=False),
                gr.update(visible=False, value=""),
                gr.update(interactive=False),
                gr.update(interactive=True),
                empty_rate_plot(app_config),
                render_text_progress(app_config, None),
                render_audio_stream(app_config, session_id=uuid.uuid4().hex),
                *generation_button_updates(running=False),
                "",
            )

        clear_btn.click(
            fn=clear_generation,
            inputs=generation_session,
            outputs=[
                prompt_audio,
                target_text,
                output_audio,
                validation_msg,
                submit_btn,
                enable_speaking_rate,
                rate_plot,
                text_progress,
                stream_audio,
                pause_btn,
                resume_btn,
                stop_btn,
                generation_session,
            ],
        )

        gr.Markdown("### Examples")
        ex = gr.Examples(
            examples=demo_examples,
            inputs=[
                prompt_audio,
                target_text,
                prompt_enhancement,
                voice_activity_detection,
                streaming_input,
                speaking_rate_control,
                enable_speaking_rate,
            ],
            outputs=[
                output_audio,
                enable_speaking_rate,
                rate_plot,
                text_progress,
                stream_audio,
                pause_btn,
                resume_btn,
                stop_btn,
                generation_session,
            ],
            fn=synthesize_fn,
            cache_examples=False,
        )

        ex.dataset.click(
            fn=lambda: (*clear_outputs(app_config), ""),
            inputs=[],
            outputs=[
                output_audio,
                rate_plot,
                text_progress,
                stream_audio,
                generation_session,
            ],
            queue=False,
        ).then(
            fn=validate_inputs,
            inputs=[prompt_audio, target_text],
            outputs=[validation_msg, submit_btn],
            queue=False,
        )

    demo.queue(default_concurrency_limit=1).launch()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--config",
        type=existing_file,
        help="Path to the config file",
        default="configs/generator.json",
    )
    parser.add_argument(
        "--app-config",
        type=existing_file,
        help="Path to the app config file",
        default="configs/app.json",
    )
    parser.add_argument(
        "--spk-rate-config",
        type=existing_file,
        help="Path to the speaking rate config file",
        default="configs/speaking_rate.json",
    )
    parser.add_argument(
        "--examples-config",
        type=existing_file,
        help="Path to the examples config file",
        default="assets/examples.json",
    )
    args = parser.parse_args()

    with open(args.config) as f:
        config = SpeechGeneratorConfig(**json.load(f))
    config.hf_token = os.environ.get("TOKEN")

    # Loading speaker encoder.
    torch.hub.load(
        config.spk_enc_repo,
        config.spk_enc_model,
        model_name=config.spk_enc_model_name,
        train_type=config.spk_enc_train_type,
        dataset=config.spk_enc_dataset,
        trust_repo=True,
        verbose=False,
    )

    with open(args.spk_rate_config) as f:
        spk_rate_config = json.load(f)

    app_config = load_app_config(args.app_config)

    with open(args.examples_config) as f:
        examples_config = json.load(f)
    demo_examples = examples_config.get("examples", [])

    speech_generator = SpeechGenerator(config, spk_rate_config)
    speaking_rate_state = SpeakingRateState(app_config.speaking_rate_default)
    generation_control = GenerationControl()
    shared_generation_state = SharedGenerationState()
    chunk_size = int(config.mimi_sr * app_config.min_chunk_sec)

    @spaces.GPU
    def synthesize_fn(
        prompt_audio_path,
        target_text,
        prompt_enhancement,
        voice_activity_detection,
        streaming_input,
        speaking_rate_control,
        enable_speaking_rate=True,
        generation_session_id="",
    ):
        control_session_id = generation_session_id or shared_generation_state.create(
            speaking_rate_control
        )
        stream_session_id = control_session_id or uuid.uuid4().hex
        stream_seq = 0

        if not prompt_audio_path or not target_text:
            speaking_rate_state.stop()
            generation_control.finish()
            shared_generation_state.finish(control_session_id)
            yield (
                gr.update(value=None, visible=False),
                gr.update(interactive=True),
                empty_rate_plot(app_config, show_target=enable_speaking_rate),
                render_text_progress(app_config, None),
                render_audio_stream(app_config, session_id=stream_session_id),
                *generation_button_updates(running=False),
                control_session_id,
            )
            return

        if shared_generation_state.is_stopped(control_session_id):
            speaking_rate_state.stop()
            generation_control.finish()
            shared_generation_state.finish(control_session_id)
            yield (
                gr.update(value=None, visible=False),
                gr.update(interactive=True),
                empty_rate_plot(app_config, show_target=enable_speaking_rate),
                render_text_progress(app_config, None),
                render_audio_stream(
                    app_config,
                    session_id=stream_session_id,
                    active=False,
                    final=True,
                ),
                *generation_button_updates(running=False),
                control_session_id,
            )
            return

        ensure_generator_on_cuda(speech_generator)
        speaking_rate_state.ensure_started(speaking_rate_control)
        speaking_rate_gen = (
            shared_generation_state.speaking_rate_values(
                control_session_id, speaking_rate_control
            )
            if enable_speaking_rate
            else None
        )
        text_metadata = build_text_progress_metadata(
            target_text,
            config=config,
            phone_to_token=speech_generator.ctx.phone_to_token,
            phonemizer=speech_generator.ctx.phonemizer,
            max_phone_tokens=config.max_phone_tokens,
        )
        rate_window_sec = (
            config.spk_rate_window_sec
            if config.spk_rate_window_sec and config.spk_rate_window_sec > 0
            else app_config.plot_window_sec
        )
        frame_sec = config.mimi_frame_ms / 1000.0
        text_progress_delay_sec = (
            app_config.audio_stream_start_delay_sec
            + config.audio_delay_frames * frame_sec
        )
        visualization = VisualizationState(
            text_metadata=text_metadata,
            app_config=app_config,
            rate_window_sec=rate_window_sec,
            frame_sec=frame_sec,
            text_progress_delay_sec=text_progress_delay_sec,
            show_target=enable_speaking_rate,
        )

        stream = speech_generator.generate_stream(
            prompt_audio_path=Path(prompt_audio_path),
            text=text_generator(target_text) if streaming_input else target_text,
            speaking_rate=speaking_rate_gen,
            enhance_prompt=prompt_enhancement,
            apply_vad=voice_activity_detection,
            return_progress=True,
            min_streaming_rtf=app_config.min_streaming_rtf,
        )

        buffer = []
        buffer_len = 0
        total_buffer = []
        stopped = False

        stream_iter = iter(stream)
        while True:
            if not shared_generation_state.wait_if_paused(control_session_id):
                stopped = True
                break
            try:
                frame, _, progress = next(stream_iter)
            except StopIteration:
                break
            if shared_generation_state.is_stopped(control_session_id):
                stopped = True
                break

            buffer.append(frame)
            total_buffer.append(frame)
            buffer_len += frame.shape[0]
            plot_update, text_update = visualization.update(progress)

            if buffer_len >= chunk_size:
                if shared_generation_state.is_stopped(control_session_id):
                    stopped = True
                    break
                audio = np.concatenate(buffer)
                stream_seq += 1
                yield (
                    gr.update(),
                    gr.update(),
                    plot_update,
                    text_update,
                    render_audio_stream(
                        app_config,
                        session_id=stream_session_id,
                        seq=stream_seq,
                        sample_rate=config.mimi_sr,
                        audio=float32_to_int16(audio),
                        active=True,
                    ),
                    *generation_button_updates(
                        running=True,
                        paused=shared_generation_state.is_paused(control_session_id),
                    ),
                    control_session_id,
                )

                buffer = []
                buffer_len = 0

        stopped = stopped or shared_generation_state.is_stopped(control_session_id)
        if stopped and hasattr(stream, "close"):
            stream.close()
        final_text = visualization.final_text()

        if buffer_len > 0 and not stopped:
            final = np.concatenate(buffer)
            nfade = min(int(config.mimi_sr * app_config.fade_out_sec), final.shape[0])
            if nfade > 0:
                fade = np.linspace(1.0, 0.0, nfade, dtype=np.float32)
                final[-nfade:] *= fade
            stream_seq += 1
            yield (
                gr.update(),
                gr.update(),
                visualization.latest_plot,
                visualization.latest_text,
                render_audio_stream(
                    app_config,
                    session_id=stream_session_id,
                    seq=stream_seq,
                    sample_rate=config.mimi_sr,
                    audio=float32_to_int16(final),
                    active=True,
                ),
                *generation_button_updates(
                    running=True,
                    paused=shared_generation_state.is_paused(control_session_id),
                ),
                control_session_id,
            )

        if len(total_buffer) > 0:
            full_audio = np.concatenate(total_buffer)
            nfade = min(
                int(config.mimi_sr * app_config.fade_out_sec), full_audio.shape[0]
            )
            if nfade > 0:
                fade = np.linspace(1.0, 0.0, nfade, dtype=np.float32)
                full_audio[-nfade:] *= fade

            file_path = f"/tmp/voxtream_{uuid.uuid4().hex}.wav"
            sf.write(file_path, float32_to_int16(full_audio), config.mimi_sr)

            speaking_rate_state.stop()
            generation_control.finish()
            shared_generation_state.finish(control_session_id)
            yield (
                gr.update(value=file_path, visible=True),
                gr.update(interactive=True),
                visualization.latest_plot,
                final_text,
                render_audio_stream(
                    app_config,
                    session_id=stream_session_id,
                    seq=stream_seq + 1,
                    sample_rate=config.mimi_sr,
                    active=False,
                    final=True,
                ),
                *generation_button_updates(running=False),
                control_session_id,
            )
        else:
            speaking_rate_state.stop()
            generation_control.finish()
            shared_generation_state.finish(control_session_id)
            yield (
                gr.update(value=None, visible=False),
                gr.update(interactive=True),
                visualization.latest_plot,
                final_text,
                render_audio_stream(
                    app_config,
                    session_id=stream_session_id,
                    seq=stream_seq + 1,
                    sample_rate=config.mimi_sr,
                    active=False,
                    final=True,
                ),
                *generation_button_updates(running=False),
                control_session_id,
            )

    demo_app(
        config,
        app_config,
        demo_examples,
        synthesize_fn,
        speaking_rate_state,
        generation_control,
        shared_generation_state,
    )


if __name__ == "__main__":
    main()