app.py

from __future__ import annotations

import os
from pathlib import Path
import sys
from uuid import uuid4

import gradio as gr

try:
    import spaces
except ModuleNotFoundError:  # Local and paid-GPU runtimes do not require the package.
    spaces = None


ROOT = Path(__file__).resolve().parent
for relative in ("vendor/chatterbox_tts_0_1_7", "src", "scripts"):
    path = ROOT / relative
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from urdu_s2s.api import DEFAULT_PRAXY_ANCHOR, result_to_api_payload, run_live_tts
from urdu_s2s.live_providers import (
    OpenAIBridgeProvider,
    OpenAICompatibleChatClient,
    neutralize_devanagari_user_address,
    normalize_devanagari_tts_pronunciation,
)
from urdu_s2s.schemas import BridgeResult, ReplyResult, SpeechToSpeechRequest
from urdu_s2s.tts_providers import ChatterboxPraxyTTSProvider


OUTPUT_DIR = ROOT / "artifacts/gradio_tts_outputs"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


def env(name: str, default: str) -> str:
    return os.environ.get(name, default)


def gpu_runtime(func):
    if spaces is None:
        return func
    return spaces.GPU(duration=int(env("S2S_ZEROGPU_DURATION", "120")))(func)


def is_devanagari(text: str) -> bool:
    return any("\u0900" <= char <= "\u097f" for char in text)


def voice_prompt_audio_path() -> Path:
    return Path(env("S2S_VOICE_PROMPT_AUDIO_PATH", str(DEFAULT_PRAXY_ANCHOR)))


def chatterbox_device() -> str:
    return env("S2S_CHATTERBOX_DEVICE", "cuda")


def chatterbox_t3_model() -> str:
    return env("S2S_CHATTERBOX_T3_MODEL", "v3")


@gpu_runtime
def synthesize_s2s(mic_audio_path: str | None, upload_audio_path: str | None, prompt_roman_urdu: str):
    audio_path = mic_audio_path or upload_audio_path
    if not audio_path:
        raise gr.Error("Please stop the recorder or upload an audio file first.")

    request_id = f"hf_{uuid4().hex}"
    input_audio = Path(audio_path)
    common = {
        "audio_path": input_audio,
        "request_id": request_id,
        "prompt_roman_urdu": prompt_roman_urdu or "",
        "whisper_model": env("S2S_WHISPER_MODEL", "small"),
        "whisper_language": env("S2S_WHISPER_LANGUAGE", "ur"),
        "whisper_device": env("S2S_WHISPER_DEVICE", "cpu"),
        "whisper_compute_type": env("S2S_WHISPER_COMPUTE_TYPE", "int8"),
    }

    result = run_live_tts(
        **common,
        output_audio_path=OUTPUT_DIR / f"{request_id}.wav",
        voice_prompt_audio_path=voice_prompt_audio_path(),
        chatterbox_device=chatterbox_device(),
        chatterbox_t3_model=chatterbox_t3_model(),
    )

    payload = result_to_api_payload(result, repo_root=ROOT)
    output_audio = payload["tts_audio_path"] if payload["tts_audio_exists"] else None
    return (
        output_audio,
        payload["asr_transcript"],
        payload["assistant_reply_urdu"],
    )


@gpu_runtime
def synthesize_tts(input_text: str):
    text = (input_text or "").strip()
    if not text:
        raise gr.Error("Please enter the text you want spoken.")

    request_id = f"hf_tts_{uuid4().hex}"
    request = SpeechToSpeechRequest(
        request_id=request_id,
        audio_path=Path(""),
        metadata={"input_mode": "text"},
    )
    reply = ReplyResult(text_urdu=text, provider="user_text", model="typed_input")

    if is_devanagari(text):
        speech_text = neutralize_devanagari_user_address(
            normalize_devanagari_tts_pronunciation(text)
        )
        bridge = BridgeResult(
            text_devanagari=speech_text,
            provider="direct_devanagari_text",
            model="user_input",
        )
    else:
        bridge = OpenAIBridgeProvider(chat_client=OpenAICompatibleChatClient()).convert(
            request,
            reply,
        )

    tts_result = ChatterboxPraxyTTSProvider(
        output_audio_path=OUTPUT_DIR / f"{request_id}.wav",
        voice_prompt_audio_path=voice_prompt_audio_path(),
        device=chatterbox_device(),
        t3_model=chatterbox_t3_model(),
    ).synthesize(request, reply, bridge)

    output_audio = str(tts_result.audio_path) if tts_result.audio_path.exists() else None
    return output_audio


with gr.Blocks(title="Urdu S2S") as demo:
    gr.Markdown("# Urdu S2S")
    gr.Markdown(
        """
        A Pakistani Urdu speech model demo for conversational voice responses. It supports
        both speech-to-speech and text-to-speech testing with the same Urdu voice style.
        """
    )
    with gr.Accordion("About this model", open=False):
        gr.Markdown(
            """
            **What it is**

            Urdu S2S is an experimental Urdu voice assistant model focused on natural
            Pakistani Urdu pronunciation, short spoken replies, and a warm female assistant
            voice. It is designed for interactive product testing rather than long-form
            narration.

            **Architecture**

            Urdu S2S uses a modular speech pipeline with dedicated components for speech
            recognition, Urdu response generation, pronunciation-aware speech preparation,
            and neural speech synthesis. The user-facing experience is a single
            speech-to-speech model demo, while the underlying design lets each part improve
            independently as stronger Urdu speech models become available.

            At a high level, incoming audio is normalized and transcribed, the dialogue
            layer produces a concise Urdu assistant response, a pronunciation layer prepares
            the response for stable speech generation, and the synthesis layer renders the
            final spoken voice.

            **What it can do**

            - Record or upload Urdu speech and receive a spoken Urdu response.
            - Generate speech directly from typed Urdu, Roman Urdu, or Devanagari text.
            - Preserve common names, numbers, dates, and English terms where possible.
            - Use Urdu-first phrasing such as “baraye meherbani” instead of overly Hindi wording.
            - Avoid assuming the user's gender in normal assistant replies.

            **Best test inputs**

            Use short, natural prompts of 2 to 10 seconds. The current demo works best for
            assistant-style requests such as greetings, clarifying questions, simple customer
            support, scheduling drafts, travel questions, and account/helpdesk style prompts.

            **Current limits**

            This is an MVP. Very noisy audio, long recordings, mixed languages, unusual names,
            and dense factual questions may still produce transcription or pronunciation errors.
            The first generation after the Space wakes up can be slower while the speech model
            loads.
            """
        )
    with gr.Tabs():
        with gr.Tab("Speech to speech"):
            with gr.Row():
                with gr.Column():
                    mic_audio = gr.Audio(
                        label="Record speech",
                        sources=["microphone"],
                        type="filepath",
                    )
                    upload_audio = gr.Audio(
                        label="Upload speech",
                        sources=["upload"],
                        type="filepath",
                    )
                    prompt = gr.Textbox(
                        label="Optional Roman Urdu hint",
                        placeholder="Example: reply like a polite Pakistani Urdu assistant",
                    )
                    run_s2s = gr.Button("Generate speech response", variant="primary")
                with gr.Column():
                    output_audio = gr.Audio(label="Response audio", type="filepath")
                    transcript = gr.Textbox(label="Transcript", rtl=True)
                    reply = gr.Textbox(label="Reply", rtl=True)

            run_s2s.click(
                fn=synthesize_s2s,
                inputs=[mic_audio, upload_audio, prompt],
                outputs=[output_audio, transcript, reply],
                api_name="s2s",
                concurrency_limit=1,
            )

        with gr.Tab("Text to speech"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Text",
                        lines=5,
                        placeholder="Type Urdu, Roman Urdu, or Devanagari text to speak.",
                    )
                    run_tts = gr.Button("Generate speech", variant="primary")
                with gr.Column():
                    tts_output_audio = gr.Audio(label="Speech audio", type="filepath")

            run_tts.click(
                fn=synthesize_tts,
                inputs=[text_input],
                outputs=[tts_output_audio],
                api_name="tts",
                concurrency_limit=1,
            )


if __name__ == "__main__":
    demo.queue(max_size=8).launch(server_name="0.0.0.0", server_port=7860)