src/speaches/gradio_app.py

from collections.abc import AsyncGenerator
from pathlib import Path
import platform

import gradio as gr
import httpx
from httpx_sse import aconnect_sse
from openai import AsyncOpenAI

from speaches.config import Config, Task
from speaches.hf_utils import PiperModel

TRANSCRIPTION_ENDPOINT = "/v1/audio/transcriptions"
TRANSLATION_ENDPOINT = "/v1/audio/translations"
TIMEOUT_SECONDS = 180
TIMEOUT = httpx.Timeout(timeout=TIMEOUT_SECONDS)

# NOTE: `gr.Request` seems to be passed in as the last positional (not keyword) argument


def base_url_from_gradio_req(request: gr.Request) -> str:
    # NOTE: `request.request.url` seems to always have a path of "/gradio_api/queue/join"
    assert request.request is not None
    return f"{request.request.url.scheme}://{request.request.url.netloc}"


def http_client_from_gradio_req(request: gr.Request, config: Config) -> httpx.AsyncClient:
    base_url = base_url_from_gradio_req(request)
    return httpx.AsyncClient(
        base_url=base_url,
        timeout=TIMEOUT,
        headers={"Authorization": f"Bearer {config.api_key}"} if config.api_key else None,
    )


def openai_client_from_gradio_req(request: gr.Request, config: Config) -> AsyncOpenAI:
    base_url = base_url_from_gradio_req(request)
    return AsyncOpenAI(base_url=f"{base_url}/v1", api_key=config.api_key if config.api_key else "cant-be-empty")


def create_gradio_demo(config: Config) -> gr.Blocks:  # noqa: C901, PLR0915
    async def whisper_handler(
        file_path: str, model: str, task: Task, temperature: float, stream: bool, request: gr.Request
    ) -> AsyncGenerator[str, None]:
        http_client = http_client_from_gradio_req(request, config)
        if task == Task.TRANSCRIBE:
            endpoint = TRANSCRIPTION_ENDPOINT
        elif task == Task.TRANSLATE:
            endpoint = TRANSLATION_ENDPOINT

        if stream:
            previous_transcription = ""
            async for transcription in streaming_audio_task(http_client, file_path, endpoint, temperature, model):
                previous_transcription += transcription
                yield previous_transcription
        else:
            yield await audio_task(http_client, file_path, endpoint, temperature, model)

    async def audio_task(
        http_client: httpx.AsyncClient, file_path: str, endpoint: str, temperature: float, model: str
    ) -> str:
        with Path(file_path).open("rb") as file:  # noqa: ASYNC230
            response = await http_client.post(
                endpoint,
                files={"file": file},
                data={
                    "model": model,
                    "response_format": "text",
                    "temperature": temperature,
                },
            )

        response.raise_for_status()
        return response.text

    async def streaming_audio_task(
        http_client: httpx.AsyncClient, file_path: str, endpoint: str, temperature: float, model: str
    ) -> AsyncGenerator[str, None]:
        with Path(file_path).open("rb") as file:  # noqa: ASYNC230
            kwargs = {
                "files": {"file": file},
                "data": {
                    "response_format": "text",
                    "temperature": temperature,
                    "model": model,
                    "stream": True,
                },
            }
            async with aconnect_sse(http_client, "POST", endpoint, **kwargs) as event_source:
                async for event in event_source.aiter_sse():
                    yield event.data

    async def update_whisper_model_dropdown(request: gr.Request) -> gr.Dropdown:
        openai_client = openai_client_from_gradio_req(request, config)
        models = (await openai_client.models.list()).data
        model_names: list[str] = [model.id for model in models]
        assert config.whisper.model in model_names
        recommended_models = {model for model in model_names if model.startswith("Systran")}
        other_models = [model for model in model_names if model not in recommended_models]
        model_names = list(recommended_models) + other_models
        return gr.Dropdown(
            choices=model_names,
            label="Model",
            value=config.whisper.model,
        )

    async def update_piper_voices_dropdown(request: gr.Request) -> gr.Dropdown:
        http_client = http_client_from_gradio_req(request, config)
        res = (await http_client.get("/v1/audio/speech/voices")).raise_for_status()
        piper_models = [PiperModel.model_validate(x) for x in res.json()]
        return gr.Dropdown(choices=[model.voice for model in piper_models], label="Voice", value=DEFAULT_VOICE)

    async def handle_audio_speech(
        text: str, voice: str, response_format: str, speed: float, sample_rate: int | None, request: gr.Request
    ) -> Path:
        openai_client = openai_client_from_gradio_req(request, config)
        res = await openai_client.audio.speech.create(
            input=text,
            model="piper",
            voice=voice,  # pyright: ignore[reportArgumentType]
            response_format=response_format,  # pyright: ignore[reportArgumentType]
            speed=speed,
            extra_body={"sample_rate": sample_rate},
        )
        audio_bytes = res.response.read()
        file_path = Path(f"audio.{response_format}")
        with file_path.open("wb") as file:  # noqa: ASYNC230
            file.write(audio_bytes)
        return file_path

    with gr.Blocks(title="Speaches Playground") as demo:
        gr.Markdown(
            "### Consider supporting the project by starring the [repository on GitHub](https://github.com/speaches-ai/speaches)."
        )
        with gr.Tab(label="Transcribe/Translate"):
            audio = gr.Audio(type="filepath")
            model_dropdown = gr.Dropdown(
                choices=[config.whisper.model],
                label="Model",
                value=config.whisper.model,
            )
            task_dropdown = gr.Dropdown(
                choices=[task.value for task in Task],
                label="Task",
                value=Task.TRANSCRIBE,
            )
            temperature_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.0)
            stream_checkbox = gr.Checkbox(label="Stream", value=True)
            button = gr.Button("Generate")

            output = gr.Textbox()

            # NOTE: the inputs order must match the `whisper_handler` signature
            button.click(
                whisper_handler, [audio, model_dropdown, task_dropdown, temperature_slider, stream_checkbox], output
            )

        with gr.Tab(label="Speech Generation"):
            if platform.machine() != "x86_64":
                from speaches.routers.speech import (
                    DEFAULT_VOICE,
                    MAX_SAMPLE_RATE,
                    MIN_SAMPLE_RATE,
                    SUPPORTED_RESPONSE_FORMATS,
                )

                text = gr.Textbox(label="Input Text")
                voice_dropdown = gr.Dropdown(
                    choices=["en_US-amy-medium"],
                    label="Voice",
                    value="en_US-amy-medium",
                    info="""
The last part of the voice name is the quality (x_low, low, medium, high).
Each quality has a different default sample rate:
- x_low: 16000 Hz
- low: 16000 Hz
- medium: 22050 Hz
- high: 22050 Hz
""",
                )
                response_fromat_dropdown = gr.Dropdown(
                    choices=SUPPORTED_RESPONSE_FORMATS,
                    label="Response Format",
                    value="wav",
                )
                speed_slider = gr.Slider(minimum=0.25, maximum=4.0, step=0.05, label="Speed", value=1.0)
                sample_rate_slider = gr.Number(
                    minimum=MIN_SAMPLE_RATE,
                    maximum=MAX_SAMPLE_RATE,
                    label="Desired Sample Rate",
                    info="""
Setting this will resample the generated audio to the desired sample rate.
You may want to set this if you are going to use voices of different qualities but want to keep the same sample rate.
Default: None (No resampling)
""",
                    value=lambda: None,
                )
                button = gr.Button("Generate Speech")
                output = gr.Audio(type="filepath")
                button.click(
                    handle_audio_speech,
                    [text, voice_dropdown, response_fromat_dropdown, speed_slider, sample_rate_slider],
                    output,
                )
                demo.load(update_piper_voices_dropdown, inputs=None, outputs=voice_dropdown)
            else:
                gr.Textbox("Speech generation is only supported on x86_64 machines.")

        demo.load(update_whisper_model_dropdown, inputs=None, outputs=model_dropdown)
    return demo