app.py

from __future__ import annotations

import io
import importlib.util
import re
import sys
import threading
import warnings
import wave
from pathlib import Path
from types import ModuleType
from typing import Iterator

import gradio as gr
import numpy as np

from example_texts import EXAMPLE_TEXTS

warnings.filterwarnings("ignore", message="`torch.nn.utils.weight_norm` is deprecated")

try:
    import torch  # type: ignore
except Exception:  # pragma: no cover
    torch = None  # type: ignore

try:
    from huggingface_hub import snapshot_download  # type: ignore
except Exception:  # pragma: no cover
    snapshot_download = None  # type: ignore


MODEL_ID = "owensong/Inflect-Nano-v1"
SAMPLE_RATE = 24_000
MAX_CHUNK_CHARS = 320
CHUNK_SILENCE_SECONDS = 0.18

_INFLECT_STATE = {
    "initialized": False,
    "device": "cpu",
    "snapshot_dir": None,
    "inference": None,
    "acoustic": None,
    "vocoder": None,
    "speakers": {},
    "acoustic_params": 0,
    "vocoder_params": 0,
}
_INFLECT_LOCK = threading.Lock()


def _detect_device() -> str:
    return "cuda" if torch and torch.cuda.is_available() else "cpu"


def get_device_info() -> str:
    if _INFLECT_STATE["initialized"]:
        return str(_INFLECT_STATE["device"])
    return _detect_device()


def _load_inference_module(snapshot_dir: Path) -> ModuleType:
    module_name = "_inflect_nano_v1_inference"
    existing = sys.modules.get(module_name)
    if existing is not None:
        return existing

    repo_root = str(snapshot_dir)
    frontend_root = str(snapshot_dir / "third_party" / "tiny_tts_frontend")
    tiny_tts_root = str(snapshot_dir / "third_party" / "tiny_tts_frontend" / "tiny_tts")
    for path in (repo_root, frontend_root):
        if path not in sys.path:
            sys.path.insert(0, path)

    if "tiny_tts" not in sys.modules:
        tiny_tts_pkg = ModuleType("tiny_tts")
        tiny_tts_pkg.__path__ = [tiny_tts_root]  # type: ignore[attr-defined]
        sys.modules["tiny_tts"] = tiny_tts_pkg

    spec = importlib.util.spec_from_file_location(module_name, snapshot_dir / "inference.py")
    if spec is None or spec.loader is None:
        raise gr.Error("Could not load Inflect-Nano inference runtime from the downloaded snapshot.")

    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


def _download_snapshot() -> Path:
    if snapshot_download is None:
        raise gr.Error("huggingface_hub is not installed. Install this app's requirements first.")

    local_dir = Path(__file__).resolve().parent / ".cache" / "inflect-nano-v1"
    return Path(
        snapshot_download(
            MODEL_ID,
            local_dir=local_dir,
            allow_patterns=[
                "inference.py",
                "inflect_nano/*.py",
                "third_party/tiny_tts_frontend/**",
                "weights/*.pt",
                "model_summary.json",
                "README.md",
            ],
        )
    )


def _ensure_nltk_data() -> None:
    try:
        import nltk  # type: ignore
    except Exception as exc:  # pragma: no cover
        raise gr.Error(f"NLTK is required by g2p_en but could not be imported: {exc}")

    nltk_dir = Path(__file__).resolve().parent / ".cache" / "nltk_data"
    nltk_dir.mkdir(parents=True, exist_ok=True)
    nltk_dir_str = str(nltk_dir)
    if nltk_dir_str not in nltk.data.path:
        nltk.data.path.insert(0, nltk_dir_str)

    resources = [
        ("averaged_perceptron_tagger_eng", "taggers/averaged_perceptron_tagger_eng"),
        ("averaged_perceptron_tagger", "taggers/averaged_perceptron_tagger"),
        ("cmudict", "corpora/cmudict"),
    ]
    for package, resource in resources:
        try:
            nltk.data.find(resource)
        except LookupError:
            print(f"Downloading NLTK resource for Inflect-Nano: {package}")
            nltk.download(package, download_dir=nltk_dir_str, quiet=True)


def _init_inflect() -> None:
    if _INFLECT_STATE["initialized"]:
        return

    with _INFLECT_LOCK:
        if _INFLECT_STATE["initialized"]:
            return

        if torch is None:
            raise gr.Error("PyTorch is not installed. Install this app's requirements first.")

        device_name = _detect_device()
        device = torch.device(device_name)
        print(f"Initializing Inflect-Nano-v1 on {device_name}")

        snapshot_dir = _download_snapshot()
        inference = _load_inference_module(snapshot_dir)
        _ensure_nltk_data()
        acoustic, speakers, acoustic_params = inference.load_acoustic(inference.DEFAULT_ACOUSTIC, device)
        vocoder, vocoder_params = inference.load_vocoder(inference.DEFAULT_VOCODER, device)

        _INFLECT_STATE.update(
            {
                "initialized": True,
                "device": device_name,
                "snapshot_dir": snapshot_dir,
                "inference": inference,
                "acoustic": acoustic,
                "vocoder": vocoder,
                "speakers": speakers,
                "acoustic_params": acoustic_params,
                "vocoder_params": vocoder_params,
            }
        )
        print(
            "Inflect-Nano-v1 initialized. "
            f"Params: acoustic={acoustic_params:,} vocoder={vocoder_params:,} total={acoustic_params + vocoder_params:,}"
        )


@torch.inference_mode() if torch is not None else lambda fn: fn
def _synthesize_inflect(
    text: str,
    length_scale: float,
    pitch_scale: float,
    energy_scale: float,
    min_duration: int,
    max_duration: int,
    smooth_predictors: bool,
) -> np.ndarray:
    inference = _INFLECT_STATE["inference"]
    acoustic = _INFLECT_STATE["acoustic"]
    vocoder = _INFLECT_STATE["vocoder"]
    speakers = _INFLECT_STATE["speakers"]
    device = torch.device(str(_INFLECT_STATE["device"]))

    phone, tone, lang = inference.text_to_tokens(text)
    phone = phone.unsqueeze(0).to(device)
    tone = tone.unsqueeze(0).to(device)
    lang = lang.unsqueeze(0).to(device)
    speaker = torch.LongTensor([int(speakers.get("mark", next(iter(speakers.values()), 0)))]).to(device)

    mel = acoustic.infer(
        phone,
        tone,
        lang,
        speaker,
        length_scale=float(length_scale),
        min_duration=int(min_duration),
        max_duration=int(max_duration),
        pitch_scale=float(pitch_scale),
        energy_scale=float(energy_scale),
        smooth_predictors=bool(smooth_predictors),
    )
    wav = vocoder(mel).squeeze().detach().cpu().numpy()
    if str(_INFLECT_STATE["device"]) == "cuda":
        torch.cuda.synchronize()
    return inference.normalize_audio(wav)


def _audio_np_to_int16(audio_np: np.ndarray) -> np.ndarray:
    audio_clipped = np.clip(audio_np, -1.0, 1.0)
    return (audio_clipped * 32767.0).astype(np.int16)


def _wav_bytes_from_int16(audio_int16: np.ndarray, sample_rate: int = SAMPLE_RATE) -> bytes:
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(audio_int16.tobytes())
    return buffer.getvalue()


def _split_long_text_part(text: str, max_chars: int) -> list[str]:
    words = text.split()
    chunks: list[str] = []
    current = ""

    for word in words:
        if len(word) > max_chars:
            if current:
                chunks.append(current)
                current = ""
            chunks.extend(word[i : i + max_chars] for i in range(0, len(word), max_chars))
        elif not current:
            current = word
        elif len(current) + 1 + len(word) <= max_chars:
            current = f"{current} {word}"
        else:
            chunks.append(current)
            current = word

    if current:
        chunks.append(current)

    return chunks


def _split_text_for_synthesis(text: str, max_chars: int = MAX_CHUNK_CHARS) -> list[str]:
    normalized = re.sub(r"\s+", " ", text.strip())
    sentence_parts = re.split(r"(?<=[.!?…])\s+", normalized)
    chunks: list[str] = []
    current = ""

    for part in (item.strip() for item in sentence_parts if item.strip()):
        subparts = [part] if len(part) <= max_chars else _split_long_text_part(part, max_chars)
        for subpart in subparts:
            if not current:
                current = subpart
            elif len(current) + 1 + len(subpart) <= max_chars:
                current = f"{current} {subpart}"
            else:
                chunks.append(current)
                current = subpart

    if current:
        chunks.append(current)

    return chunks


def inflect_tts(
    text: str,
    length_scale: float,
    pitch_scale: float,
    energy_scale: float,
    min_duration: int,
    max_duration: int,
    smooth_predictors: bool,
) -> Iterator[bytes]:
    text = (text or "").strip()
    if not text:
        raise gr.Error("Please enter text to synthesize.")
    if min_duration > max_duration:
        raise gr.Error("Minimum duration must be less than or equal to maximum duration.")

    _init_inflect()

    try:
        chunks = _split_text_for_synthesis(text)
        if not chunks:
            raise gr.Error("No valid text chunks found.")

        trailing_silence = np.zeros(int(SAMPLE_RATE * CHUNK_SILENCE_SECONDS), dtype=np.float32)
        for index, chunk in enumerate(chunks, start=1):
            print(f"Inflect-Nano chunk {index}/{len(chunks)}: {len(chunk)} chars")
            audio = _synthesize_inflect(
                chunk,
                length_scale,
                pitch_scale,
                energy_scale,
                min_duration,
                max_duration,
                smooth_predictors,
            )
            if index < len(chunks):
                audio = np.concatenate([audio, trailing_silence], axis=0)
            yield _wav_bytes_from_int16(_audio_np_to_int16(audio), SAMPLE_RATE)
    except gr.Error:
        raise
    except Exception as exc:
        raise gr.Error(f"Error during speech generation: {str(exc)[:200]}...")


def warmup_inflect(
    default_length_scale: float,
    default_pitch_scale: float,
    default_energy_scale: float,
    default_min_duration: int,
    default_max_duration: int,
    default_smooth_predictors: bool,
) -> None:
    try:
        for _chunk in inflect_tts(
            "Warm up.",
            default_length_scale,
            default_pitch_scale,
            default_energy_scale,
            default_min_duration,
            default_max_duration,
            default_smooth_predictors,
        ):
            pass
    except Exception as exc:
        print(f"Inflect-Nano warm-up skipped: {exc}")


with gr.Blocks() as demo:
    gr.HTML(
        f"<h1 style='text-align: center;'>Inflect-Nano-v1</h1>"
        f"<p style='text-align: center;'>Powered by Inflect-Nano-v1 on {get_device_info().upper()} | "
        "English single-voice 24 kHz TTS</p>"
    )

    with gr.Row(variant="panel"):
        length_scale = gr.Slider(
            minimum=0.85,
            maximum=1.20,
            value=1.00,
            step=0.01,
            label="Length Scale",
            info="Adjusts predicted phoneme durations.",
        )
        pitch_scale = gr.Slider(
            minimum=0.85,
            maximum=1.15,
            value=1.00,
            step=0.01,
            label="Pitch Scale",
            info="Scales predicted pitch.",
        )
        energy_scale = gr.Slider(
            minimum=0.85,
            maximum=1.15,
            value=1.00,
            step=0.01,
            label="Energy Scale",
            info="Scales predicted vocal energy.",
        )

    text_input = gr.Textbox(
        label="Input Text",
        placeholder="Enter the text you want to convert to speech here...",
        value="Wait, are you actually being for real now? I can't believe it!",
        lines=5,
        max_lines=8,
    )

    gr.Examples(
        examples=[[text] for text in EXAMPLE_TEXTS],
        inputs=[text_input],
        label="Examples",
    )

    with gr.Accordion("Advanced Options", open=False):
        with gr.Row():
            min_duration = gr.Slider(
                minimum=1,
                maximum=20,
                value=1,
                step=1,
                label="Minimum Token Duration",
                info="Lower bound for predicted token duration frames.",
            )
            max_duration = gr.Slider(
                minimum=20,
                maximum=120,
                value=80,
                step=1,
                label="Maximum Token Duration",
                info="Upper bound for predicted token duration frames.",
            )
            smooth_predictors = gr.Checkbox(
                value=False,
                label="Smooth Predictors",
                info="Applies light smoothing to energy, brightness, and pitch predictors.",
            )

    generate_btn = gr.Button(
        "Generate Speech",
        variant="primary",
    )

    audio_output = gr.Audio(
        label="Generated Speech",
        streaming=True,
        autoplay=True,
        buttons=["download"],
    )

    generate_inputs = [
        text_input,
        length_scale,
        pitch_scale,
        energy_scale,
        min_duration,
        max_duration,
        smooth_predictors,
    ]

    generate_btn.click(
        fn=inflect_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech",
        concurrency_limit=1,
        concurrency_id="inflect_nano_tts",
    )

    text_input.submit(
        fn=inflect_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech_enter",
        concurrency_limit=1,
        concurrency_id="inflect_nano_tts",
    )


if __name__ == "__main__":
    warmup_inflect(
        float(length_scale.value),
        float(pitch_scale.value),
        float(energy_scale.value),
        int(min_duration.value),
        int(max_duration.value),
        bool(smooth_predictors.value),
    )
    demo.queue(default_concurrency_limit=1).launch(debug=False, theme="Nymbo/Nymbo_Theme")