core/modal_infra.py

"""
core/modal_infra.py — ALL Modal GPU functions live here.

app.py and the other core/ wrappers call these remote functions.
Nothing outside this file should import torch, transformers, parler_tts,
faster_whisper, or any other ML library directly.

GPU tier and container settings are read from model_config.get_compute()
so that changing MODAL_GPU or MODAL_MIN_CONTAINERS in model_config.py
propagates here automatically.
"""

import base64

import modal
from core.model_config import get_compute, get_config
from core.prompts import SCENE_SENTINEL

_compute = get_compute()

app = modal.App("rupkotha")

# Vision image: Ollama serves every stack's vision model behind one uniform API
# (MiniCPM-V for Stack A, Gemma 3 for Stacks B/C). Switching stacks only changes
# the model tag passed in from model_config.py — never this code.
_ollama_image = (
    modal.Image.debian_slim(python_version="3.12")
    # zstd is required by the Ollama install script to extract its release archive.
    .apt_install("curl", "zstd")
    .run_commands("curl -fsSL https://ollama.com/install.sh | sh")
    .pip_install("ollama>=0.3")
    # This module imports core.model_config at load; current Modal no longer
    # auto-mounts the project, so make `core` importable in the container.
    .add_local_python_source("core")
)

# Pulled Ollama models persist here across containers, so a model is downloaded
# at most once per stack (not on every cold start).
_ollama_volume = modal.Volume.from_name("rupkotha-ollama", create_if_missing=True)

# HuggingFace weights (TTS/STT) persist here so they download at most once.
_hf_volume = modal.Volume.from_name("rupkotha-hf", create_if_missing=True)
_HF_CACHE = "/root/.cache/huggingface"

# HF auth for gated repos (e.g. ai4bharat/indic-parler-tts). Reuses the existing
# 'algaeguard-secrets' Modal secret, which provides HF_TOKEN — huggingface_hub /
# transformers read HF_TOKEN from the environment automatically.
_hf_secrets = [modal.Secret.from_name("algaeguard-secrets")]

# Warm-container model caches for the TTS/STT singletons. Populated lazily inside
# their respective functions; reused across calls in the same container.
_indic_parler: dict = {}
_voxcpm: dict = {}
_whisper: dict = {}
_indictrans: dict = {}
_indictts: dict = {}

# ML image: used by the (still-stubbed) TTS/STT functions — transformers,
# faster-whisper, parler-tts. Kept separate from the vision image so each
# function pulls only what it needs and cold-starts faster.
_ml_image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git")  # needed for the git-based parler-tts install below
    .pip_install(
        "torch>=2.2",
        "transformers>=4.40",
        "faster-whisper>=1.0",
        "voxcpm",  # English TTS — VoxCPM2 (needs Python < 3.13; 3.12 here is fine)
        "indictranstoolkit",  # IndicTrans2 pre/post-processing for the Bengali pivot
        "soundfile",
        "numpy",
    )
    .run_commands(
        "pip install git+https://github.com/huggingface/parler-tts.git"
    )
    .add_local_python_source("core")  # see note on _ollama_image above
)

# AI4Bharat Indic-TTS (FastPitch + HiFi-GAN) image. AI4Bharat's repo isn't a Python
# package — inference just uses a Coqui `Synthesizer` over their checkpoints. We use
# the maintained `coqui-tts` fork (Synthesizer API unchanged) and skip their heavy
# Indic text-normalization/denoiser layer (fragile pinned deps; unneeded for our
# clean Bengali-script input). Python 3.10 for best coqui-tts/checkpoint compat.
_indictts_image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install("wget", "unzip", "libsndfile1", "espeak-ng")
    # coqui-tts 0.27 imports transformers.pytorch_utils.isin_mps_friendly, removed in
    # transformers 5.x — pin to 4.x so `import TTS` (pulls XTTS/tortoise) works.
    .pip_install("coqui-tts[codec]", "transformers<5", "torch", "torchaudio", "soundfile", "numpy")
    .add_local_python_source("core")
)
# Persists the ~1.5 GB Bengali checkpoint zip (downloaded + unzipped once).
_indictts_volume = modal.Volume.from_name("rupkotha-indictts", create_if_missing=True)

_gpu = _compute["gpu"]
_min_containers = _compute["min_containers"]


def _ensure_ollama_server() -> None:
    """Start `ollama serve` in the background if it isn't already responding.

    Idempotent: safe to call on every invocation. Runs only inside the Modal
    container (never in the local Gradio process)."""
    import subprocess
    import time
    import urllib.request

    def _ready() -> bool:
        try:
            urllib.request.urlopen("http://127.0.0.1:11434/api/version", timeout=1)
            return True
        except Exception:
            return False

    if _ready():
        return
    subprocess.Popen(
        ["ollama", "serve"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    for _ in range(120):  # up to ~60s for the server to come up
        if _ready():
            return
        time.sleep(0.5)
    raise RuntimeError("ollama server did not become ready in time")


# ─────────────────────────────────────────────────────────────────────────────
# Vision + story generation
# ─────────────────────────────────────────────────────────────────────────────

@app.function(
    gpu=_gpu,
    image=_ollama_image,
    volumes={"/root/.ollama": _ollama_volume},
    timeout=900,  # generous: covers a first-time model pull on a cold container
    min_containers=_min_containers,
)
def run_vision_story(
    image_bytes_list: list[bytes],
    prompt: str,
    model_id: str,
    options: dict | None = None,
    describe_prompt: str | None = None,
) -> str:
    """Generate a bedtime story from images + prompt via Ollama.

    Model-agnostic: the call serves whatever Ollama tag `model_id` names — Stack A
    uses MiniCPM-V 4.5. (Bengali is served separately by the fine-tuned model via
    finetune/serve_vllm.py; this Ollama path handles English.)

    Args:
        image_bytes_list: Raw bytes of each uploaded image (1–4 images).
        prompt: Full story-generation prompt (built by core/prompts.py).
        model_id: Ollama model tag from get_config().vision_model
            (e.g. 'openbmb/minicpm-v4.5').
        options: Ollama decoding params (temperature, top_p, repeat_penalty, …).
            Per-language profiles come from model_config.get_vision_options().
        describe_prompt: Lever C (two-pass). When set, the model first DESCRIBES
            the image(s) in English using this prompt, then narrates the story
            text-only — `prompt` must contain SCENE_SENTINEL, which is replaced
            with that description. Keeps perception and Bengali prose separate.

    Returns:
        Generated story text, or '' on failure.
    """
    import time

    import ollama

    _ensure_ollama_server()

    # Pull on first use; the Volume keeps the weights warm for later calls.
    # NOTE: the Stack A tag 'openbmb/minicpm-v4.5' is confirmed valid in the
    # Ollama registry.
    listed = ollama.list()
    models = getattr(listed, "models", None) or listed.get("models", []) or []
    local_tags = {(getattr(m, "model", None) or m.get("model", "")) for m in models}
    if model_id not in local_tags and f"{model_id}:latest" not in local_tags:
        print(f"[run_vision_story] pulling {model_id} (first use, may take minutes) ...")
        ollama.pull(model_id)
        _ollama_volume.commit()
        print(f"[run_vision_story] pull complete for {model_id}")

    chat_options = options or {"temperature": 0.8}

    def _chat_once(prompt_text: str, images: list[bytes]) -> tuple[str, object]:
        # The ollama client base64-encodes raw image bytes itself.
        kwargs = dict(
            model=model_id,
            messages=[
                {"role": "user", "content": prompt_text, "images": list(images or [])}
            ],
            options=chat_options,
        )
        # think=False disables MiniCPM-V's reasoning mode, which can otherwise
        # consume the whole token budget (done_reason='length') and leave the
        # answer 'content' empty. Fall back for older ollama clients.
        try:
            resp = ollama.chat(**kwargs, think=False)
        except TypeError:
            resp = ollama.chat(**kwargs)
        # ollama returns a typed ChatResponse (subscriptable) or a plain dict.
        try:
            c = resp["message"]["content"]
        except Exception:
            c = getattr(getattr(resp, "message", None), "content", "") or ""
        return (c or "").strip(), resp

    def _chat_retry(prompt_text: str, images: list[bytes]) -> str:
        # A freshly started server sometimes returns empty on the very first call
        # while the model finishes loading into VRAM — retry a couple of times.
        text, resp = _chat_once(prompt_text, images)
        for attempt in range(2):
            if text:
                break
            print(f"[run_vision_story] empty content (attempt {attempt + 1}); resp={repr(resp)[:200]}")
            time.sleep(2)
            text, resp = _chat_once(prompt_text, images)
        return text

    if describe_prompt:
        # Lever C, pass 1: describe the image(s) in English (the model's strength).
        description = _chat_retry(describe_prompt, image_bytes_list)
        print(f"[run_vision_story] scene description: {description[:200]}")
        # Pass 2: narrate from the description, text-only (no image attached).
        story_prompt = prompt.replace(SCENE_SENTINEL, description)
        return _chat_retry(story_prompt, [])

    # Single pass: image(s) + prompt together.
    return _chat_retry(prompt, image_bytes_list)


def generate_story_remote(
    images_b64: list[str],
    prompt: str,
    options: dict | None = None,
    describe_prompt: str | None = None,
) -> str:
    """Plain-callable entry point used by core/vision_story.py.

    Decodes the base64 images, reads the active vision model tag from
    get_config() internally (model names never leave model_config.py), and
    dispatches to the deployed Modal `run_vision_story` function. `options`
    carries the per-language decoding profile from get_vision_options();
    `describe_prompt` enables two-pass (Lever C) for Bengali.

    Requires the Modal app to be deployed (`modal deploy core/modal_infra.py`)
    and Modal credentials available to the Gradio process. May raise — callers
    in core/vision_story.py wrap this in try/except and fall back to a friendly
    bedtime message, so the app stays runnable even if Modal is unreachable.
    """
    image_bytes_list = [base64.b64decode(b) for b in (images_b64 or [])]
    model_id = get_config().vision_model
    fn = modal.Function.from_name("rupkotha", "run_vision_story")
    return fn.remote(image_bytes_list, prompt, model_id, options, describe_prompt)


def generate_story_ft_remote(images_b64: list[str], prompt: str) -> str:
    """Plain-callable entry point for the Bengali-fine-tuned model, deployed
    separately as the `rupkotha-ft-serve` app (finetune/serve_vllm.py — merged
    LoRA served via vLLM). Used by core/vision_story.py only when
    model_config.FINETUNED_VISION_MODEL is set. Mirrors generate_story_remote's
    contract; may raise — callers wrap in try/except. Kept here (not imported from
    finetune/) so core/ stays independent of the training package."""
    image_bytes_list = [base64.b64decode(b) for b in (images_b64 or [])]
    fn = modal.Function.from_name("rupkotha-ft-serve", "run_vision_story_ft")
    return fn.remote(image_bytes_list, prompt)


# ─────────────────────────────────────────────────────────────────────────────
# Translation — IndicTrans2 (English → Bengali "pivot" path)
# ─────────────────────────────────────────────────────────────────────────────

@app.function(
    gpu=_gpu,
    image=_ml_image,
    volumes={_HF_CACHE: _hf_volume},
    secrets=_hf_secrets,
    timeout=600,  # generous: covers a first-time weight download on a cold start
    min_containers=_min_containers,
)
def run_translate(text: str, src_lang: str, tgt_lang: str, model_repo: str) -> str:
    """Translate text with IndicTrans2 (e.g. English → Bengali).

    Args:
        text: Source text (may be multiple sentences).
        src_lang / tgt_lang: IndicTrans2 FLORES codes ('eng_Latn', 'ben_Beng').
        model_repo: HF repo from get_config-side TRANSLATION_MODEL.

    Returns:
        Translated text, or '' on failure.
    """
    import re

    import torch
    from IndicTransToolkit import IndicProcessor
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

    if not (text or "").strip():
        return ""

    # Lazy singleton — load once per warm container.
    if "model" not in _indictrans:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        tokenizer = AutoTokenizer.from_pretrained(model_repo, trust_remote_code=True)
        model = AutoModelForSeq2SeqLM.from_pretrained(
            model_repo, trust_remote_code=True
        ).to(device)
        _indictrans.update(
            model=model, tokenizer=tokenizer, ip=IndicProcessor(inference=True), device=device
        )
        _hf_volume.commit()

    model = _indictrans["model"]
    tokenizer = _indictrans["tokenizer"]
    ip = _indictrans["ip"]
    device = _indictrans["device"]

    # IndicTrans2 translates sentence-by-sentence; split the story first.
    sentences = [s.strip() for s in re.split(r"(?<=[.!?।])\s+", text.strip()) if s.strip()]
    if not sentences:
        return ""

    batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=tgt_lang)
    enc = tokenizer(
        batch, padding="longest", truncation=True, max_length=256, return_tensors="pt"
    ).to(device)
    with torch.inference_mode():
        out = model.generate(**enc, num_beams=5, num_return_sequences=1, max_length=256)
    decoded = tokenizer.batch_decode(
        out, skip_special_tokens=True, clean_up_tokenization_spaces=True
    )
    translated = ip.postprocess_batch(decoded, lang=tgt_lang)
    return " ".join(t.strip() for t in translated if t and t.strip())


def translate_remote(text: str, src_language: str, tgt_language: str) -> str:
    """Plain-callable entry point used by core/vision_story.py for the Bengali
    translation pivot. Resolves the model + FLORES codes from model_config and
    dispatches to the deployed Modal `run_translate`. May raise — caller falls back.
    """
    from core.model_config import TRANSLATION_MODEL, get_indictrans_code

    src = get_indictrans_code(src_language)
    tgt = get_indictrans_code(tgt_language)
    fn = modal.Function.from_name("rupkotha", "run_translate")
    return fn.remote(text, src, tgt, TRANSLATION_MODEL)


# ─────────────────────────────────────────────────────────────────────────────
# Bengali TTS — Indic Parler-TTS
# ─────────────────────────────────────────────────────────────────────────────

@app.function(
    gpu=_gpu,
    image=_ml_image,
    volumes={_HF_CACHE: _hf_volume},
    secrets=_hf_secrets,
    timeout=600,  # generous: covers a cold-start model download + load
    min_containers=_min_containers,
)
def run_tts_bengali(text: str, caption: str, model_repo: str) -> bytes:
    """Synthesise Bengali speech using Indic Parler-TTS.

    Args:
        text: Story text in Bengali.
        caption: Voice-description prompt that controls the speaker persona.
        model_repo: HuggingFace repo ID from get_tts_repo() (e.g.
            'ai4bharat/indic-parler-tts').

    Returns:
        WAV audio as bytes, or b'' on failure.
    """
    import io
    import re

    import numpy as np
    import soundfile as sf
    import torch
    from parler_tts import ParlerTTSForConditionalGeneration
    from transformers import AutoTokenizer, set_seed

    # Lazy singleton — load once per warm container, reuse across calls.
    if "model" not in _indic_parler:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = ParlerTTSForConditionalGeneration.from_pretrained(model_repo).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_repo)
        desc_tokenizer = AutoTokenizer.from_pretrained(
            model.config.text_encoder._name_or_path
        )
        _indic_parler.update(
            model=model,
            tokenizer=tokenizer,
            desc_tokenizer=desc_tokenizer,
            device=device,
            sampling_rate=model.config.sampling_rate,
        )
        _hf_volume.commit()  # persist downloaded weights for the next cold start

    model = _indic_parler["model"]
    tokenizer = _indic_parler["tokenizer"]
    desc_tokenizer = _indic_parler["desc_tokenizer"]
    device = _indic_parler["device"]

    sr = _indic_parler["sampling_rate"]

    # Indic Parler-TTS caps each generation at ~30s of audio, so synthesising a
    # whole 150-word story in one pass truncates it (audio stops mid-script) and
    # rushes the prosody. Instead render it sentence-by-sentence and stitch the
    # segments with a short silence — no cut-offs, and a natural bedtime pause at
    # each punctuation mark.
    def _chunk(t: str, max_chars: int = 220) -> list[str]:
        # Split on Bengali daari (।) and ? ! … . — keep the delimiter attached,
        # then pack consecutive sentences up to max_chars so chunks stay well
        # under the ~30s generation limit.
        parts = [p for p in re.split(r"(?<=[।!?.…])\s+", t.strip()) if p]
        chunks, cur = [], ""
        for p in parts:
            if cur and len(cur) + len(p) + 1 > max_chars:
                chunks.append(cur)
                cur = p
            else:
                cur = f"{cur} {p}".strip()
        if cur:
            chunks.append(cur)
        return chunks or [t.strip()]

    desc = desc_tokenizer(caption, return_tensors="pt").to(device)
    gap = np.zeros(int(sr * 0.35), dtype=np.float32)  # ~0.35s pause between sentences
    # Sample instead of greedy decoding: greedy makes Parler sound flat/monotone.
    # do_sample + temperature gives more lively, natural prosody (the closest lever
    # Bengali has — it has no emotion-prompt support, unlike VoxCPM2 for English).
    # Fixed seed keeps a given story's audio reproducible across runs.
    set_seed(0)
    segments: list = []
    for chunk in _chunk(text):
        prompt = tokenizer(chunk, return_tensors="pt").to(device)
        with torch.no_grad():
            generation = model.generate(
                input_ids=desc.input_ids,
                attention_mask=desc.attention_mask,
                prompt_input_ids=prompt.input_ids,
                prompt_attention_mask=prompt.attention_mask,
                do_sample=True,
                temperature=1.0,
            )
        seg = generation.cpu().numpy().squeeze().astype(np.float32)
        if seg.size == 0:
            continue
        segments.append(seg)
        segments.append(gap)

    if not segments:
        return b""
    audio = np.concatenate(segments[:-1])  # drop the trailing gap
    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV")
    return buf.getvalue()


@app.function(
    gpu=_gpu,
    image=_indictts_image,
    volumes={"/models": _indictts_volume},
    timeout=1800,  # first call downloads + unzips the ~1.5 GB checkpoint
    min_containers=_min_containers,
)
def run_tts_indic_ai4bharat(text: str, checkpoint_url: str) -> bytes:
    """Synthesise Bengali speech with AI4Bharat Indic-TTS (FastPitch + HiFi-GAN).

    No reference clip — a dedicated, MOS-tuned Bengali acoustic model with a fixed
    voice. The language checkpoint zip is downloaded + unzipped once into a volume.

    Returns WAV bytes (model sample rate), or b'' on failure.
    """
    import glob
    import io
    import os
    import re
    import subprocess

    import numpy as np
    import soundfile as sf
    import torch

    ckpt_dir = "/models/indic_tts_bn"
    # One-time download + unzip into the persistent volume.
    if not glob.glob(f"{ckpt_dir}/**/fastpitch/best_model.pth", recursive=True):
        os.makedirs(ckpt_dir, exist_ok=True)
        zip_path = "/tmp/indictts_bn.zip"
        subprocess.run(["wget", "-q", "-O", zip_path, checkpoint_url], check=True)
        subprocess.run(["unzip", "-o", "-q", zip_path, "-d", ckpt_dir], check=True)
        os.remove(zip_path)
        _indictts_volume.commit()

    def _cfg_near(ckpt_path: str) -> str | None:
        d = os.path.dirname(ckpt_path)
        for cand in (os.path.join(d, "config.json"),
                     os.path.join(os.path.dirname(d), "config.json")):
            if os.path.exists(cand):
                return cand
        hits = glob.glob(os.path.join(os.path.dirname(d), "**", "config.json"), recursive=True)
        return hits[0] if hits else None

    # Lazy singleton — build the Coqui Synthesizer once per warm container.
    if "syn" not in _indictts:
        from TTS.utils.synthesizer import Synthesizer

        fp = sorted(glob.glob(f"{ckpt_dir}/**/fastpitch/best_model.pth", recursive=True))
        voc = sorted(glob.glob(f"{ckpt_dir}/**/hifigan/best_model.pth", recursive=True))
        if not fp or not voc:
            print("[indic_tts] checkpoints not found:", glob.glob(f"{ckpt_dir}/**", recursive=True)[:20])
            return b""
        import json

        spk_file = os.path.join(os.path.dirname(fp[0]), "speakers.pth")
        fp_cfg_path = _cfg_near(fp[0])
        # The config bakes a RELATIVE speakers_file path from AI4Bharat's training
        # tree (resolved against CWD → FileNotFoundError). Rewrite it to our
        # absolute path so coqui loads the right speaker map.
        if fp_cfg_path and os.path.exists(spk_file):
            with open(fp_cfg_path) as f:
                cfg_json = json.load(f)
            cfg_json["speakers_file"] = spk_file
            if isinstance(cfg_json.get("model_args"), dict):
                cfg_json["model_args"]["speakers_file"] = spk_file
            fp_cfg_path = "/tmp/fastpitch_config_patched.json"
            with open(fp_cfg_path, "w") as f:
                json.dump(cfg_json, f)
        syn = Synthesizer(
            tts_checkpoint=fp[0],
            tts_config_path=fp_cfg_path,
            tts_speakers_file=spk_file if os.path.exists(spk_file) else None,
            vocoder_checkpoint=voc[0],
            vocoder_config=_cfg_near(voc[0]),
            use_cuda=torch.cuda.is_available(),
        )
        _indictts["syn"] = syn
        # Multi-speaker model (trained on male+female): prefer the female voice for
        # the grandmother persona. AI4Bharat names them literally "male"/"female".
        speaker = None
        try:
            names = list(syn.tts_model.speaker_manager.name_to_id.keys())
            speaker = next((s for s in names if "fem" in s.lower()), names[0]) if names else None
        except Exception:  # noqa: BLE001 — single-speaker model, no manager
            speaker = None
        _indictts["speaker"] = speaker
        print("[indic_tts] loaded; speakers available:",
              locals().get("names", "?"), "| using:", speaker)

    syn = _indictts["syn"]
    speaker = _indictts.get("speaker")
    sr = syn.output_sample_rate

    def _chunk(t: str, max_chars: int = 220, min_chars: int = 8) -> list[str]:
        parts = [p for p in re.split(r"(?<=[।!?.…])\s+", t.strip()) if p.strip()]
        chunks, cur = [], ""
        for p in parts:
            if cur and len(cur) + len(p) + 1 > max_chars:
                chunks.append(cur)
                cur = p
            else:
                cur = f"{cur} {p}".strip()
        if cur:
            chunks.append(cur)
        # Merge too-short fragments (a lone quote/word) into a neighbour — FastPitch's
        # conv kernel errors when a chunk is shorter than the kernel size.
        merged: list = []
        for c in chunks:
            if merged and len(c.strip()) < min_chars:
                merged[-1] = (merged[-1] + " " + c).strip()
            else:
                merged.append(c)
        if len(merged) > 1 and len(merged[0].strip()) < min_chars:
            merged[1] = (merged[0] + " " + merged[1]).strip()
            merged = merged[1:]
        return [c for c in merged if c.strip()] or [t.strip()]

    gap = np.zeros(int(sr * 0.35), dtype=np.float32)
    segments: list = []
    for chunk in _chunk(text):
        if len(chunk.strip()) < 2:  # never feed FastPitch a 1-char chunk
            continue
        kwargs = {"speaker_name": speaker} if speaker else {}
        try:
            wav = syn.tts(chunk, **kwargs)
        except Exception as e:  # noqa: BLE001 — skip a bad chunk, don't fail the story
            print(f"[indic_tts] chunk skipped ({e}): {chunk[:40]!r}", flush=True)
            continue
        seg = np.asarray(wav, dtype=np.float32).squeeze()
        if seg.size == 0:
            continue
        segments.append(seg)
        segments.append(gap)

    if not segments:
        return b""
    audio = np.concatenate(segments[:-1])
    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV")
    return buf.getvalue()


def synthesize_bengali_remote(text: str, caption: str) -> bytes:
    """Plain-callable entry point used by core/tts.py for the Bengali path.

    Reads the active Bengali TTS backend from get_config() and dispatches to the
    right deployed Modal function. May raise — core/tts.py wraps this and falls
    back to text-only so audio failure never breaks the app.

    - 'indic_tts': AI4Bharat Indic-TTS (FastPitch + HiFi-GAN), no reference clip.
    - 'indic_parler': description-controlled Indic Parler-TTS.
    """
    from core.model_config import get_tts_repo

    cfg = get_config()
    backend = cfg.tts_bn_backend

    if backend == "indic_tts":
        # AI4Bharat Indic-TTS (FastPitch + HiFi-GAN) — no caption/reference; the
        # 'repo' here is the checkpoint-zip URL.
        fn = modal.Function.from_name("rupkotha", "run_tts_indic_ai4bharat")
        return fn.remote(text, get_tts_repo("indic_tts"))

    fn = modal.Function.from_name("rupkotha", "run_tts_bengali")
    return fn.remote(text, caption, get_tts_repo("indic_parler"))


# ─────────────────────────────────────────────────────────────────────────────
# English TTS — VoxCPM2
# ─────────────────────────────────────────────────────────────────────────────

@app.function(
    gpu=_gpu,
    image=_ml_image,
    volumes={_HF_CACHE: _hf_volume},
    secrets=_hf_secrets,
    timeout=900,  # generous: covers a cold-start model download + load
    min_containers=_min_containers,
)
def run_tts_english(text: str, voice_prompt: str, model_repo: str) -> bytes:
    """Synthesise English speech using VoxCPM2 Voice Design.

    The persona is supplied via Voice Design: VoxCPM2 reads a parenthetical
    description at the start of the text and generates a matching novel voice
    (no reference audio needed).

    Args:
        text: Story text in English.
        voice_prompt: Voice Design persona description (without parentheses).
        model_repo: HuggingFace repo ID from get_tts_repo() (e.g. 'openbmb/VoxCPM2').

    Returns:
        WAV audio as bytes, or b'' on failure.
    """
    import io

    import soundfile as sf
    from voxcpm import VoxCPM

    # Lazy singleton — load once per warm container, reuse across calls.
    if "model" not in _voxcpm:
        model = VoxCPM.from_pretrained(model_repo, load_denoiser=False)
        _voxcpm.update(model=model, sampling_rate=model.tts_model.sample_rate)
        _hf_volume.commit()  # persist downloaded weights for the next cold start

    model = _voxcpm["model"]
    # Voice Design: the persona goes in parentheses at the start of the text.
    design_text = f"({voice_prompt}){text}"
    audio = model.generate(
        text=design_text,
        cfg_value=2.0,
        inference_timesteps=10,
    )

    buf = io.BytesIO()
    sf.write(buf, audio, _voxcpm["sampling_rate"], format="WAV")
    return buf.getvalue()


def synthesize_english_remote(text: str, voice_prompt: str) -> bytes:
    """Plain-callable entry point used by core/tts.py for the English path.

    Reads the active English TTS backend from get_config(), resolves it to a
    repo ID via get_tts_repo(), and dispatches to the deployed Modal
    `run_tts_english` function. May raise — core/tts.py wraps this and falls
    back to text-only so audio failure never breaks the app.
    """
    from core.model_config import get_tts_repo

    repo = get_tts_repo(get_config().tts_en_backend)
    fn = modal.Function.from_name("rupkotha", "run_tts_english")
    return fn.remote(text, voice_prompt, repo)


# ─────────────────────────────────────────────────────────────────────────────
# Speech-to-text — faster-whisper
# ─────────────────────────────────────────────────────────────────────────────

def _stt_faster_whisper(audio_bytes: bytes, model_size: str, language: str) -> str:
    """Transcribe with faster-whisper (used for size tags like 'large-v3')."""
    import io

    import torch
    from faster_whisper import WhisperModel

    key = ("fw", model_size)
    if key not in _whisper:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        compute_type = "float16" if device == "cuda" else "int8"
        _whisper[key] = WhisperModel(
            model_size, device=device, compute_type=compute_type, download_root=_HF_CACHE
        )
        _hf_volume.commit()  # persist the downloaded model for the next cold start

    model = _whisper[key]
    segments, _ = model.transcribe(
        io.BytesIO(audio_bytes), language=language, vad_filter=True
    )
    return "".join(seg.text for seg in segments).strip()


def _stt_transformers(audio_bytes: bytes, model_repo: str) -> str:
    """Transcribe with a HF transformers ASR checkpoint (e.g. a Bengali-specific
    fine-tuned Whisper). The pipeline ffmpeg-decodes raw bytes and resamples."""
    import torch
    from transformers import pipeline

    key = ("hf", model_repo)
    if key not in _whisper:
        device = 0 if torch.cuda.is_available() else -1
        _whisper[key] = pipeline(
            "automatic-speech-recognition", model=model_repo, device=device
        )
        _hf_volume.commit()

    result = _whisper[key](audio_bytes)
    return (result.get("text") or "").strip()


@app.function(
    gpu=_gpu,
    image=_ml_image,
    volumes={_HF_CACHE: _hf_volume},
    secrets=_hf_secrets,
    timeout=600,  # generous: covers a cold-start model download + load
    min_containers=_min_containers,
)
def run_stt(audio_bytes: bytes, language: str, model: str) -> str:
    """Transcribe audio to text.

    Two backends, chosen by the model identifier:
      - a faster-whisper size tag (e.g. 'large-v3', no '/') → faster-whisper
      - a HuggingFace repo (e.g. 'bangla-asr/whisper-medium-bn', has '/')
        → transformers ASR pipeline (Bengali-specific models)

    Args:
        audio_bytes: Raw audio bytes (any format ffmpeg accepts).
        language: 'en' or 'bn'.
        model: stt_model (EN) or stt_bn_model (BN) from get_config().

    Returns:
        Transcribed text, or '' on failure (caller falls back to typed input).
    """
    if not audio_bytes:
        return ""
    try:
        if "/" in model:
            return _stt_transformers(audio_bytes, model)
        return _stt_faster_whisper(audio_bytes, model, language)
    except Exception as e:  # noqa: BLE001 — never raise; caller falls back to text
        print(f"[modal_infra] run_stt failed: {e}")
        return ""


def transcribe_remote(audio_bytes: bytes, language: str, model: str) -> str:
    """Plain-callable entry point used by core/stt.py.

    Dispatches to the deployed Modal `run_stt` function. May raise — core/stt.py
    wraps this and returns '' so the caller falls back to typed input.
    """
    fn = modal.Function.from_name("rupkotha", "run_stt")
    return fn.remote(audio_bytes, language, model)