batch_inference.py

#!/usr/bin/env python3
"""Batch offline TTS inference with pre-computed speaker embeddings.

Processes a manifest of (speaker_id, text) pairs through the vLLM-Omni
two-stage pipeline, using pre-computed embeddings to skip the ECAPA-TDNN
speaker encoder entirely.

Usage:
    python batch_inference.py \\
        --manifest manifest.json \\
        --profiles-dir voice_profiles \\
        --output-dir output_audio \\
        --model Qwen/Qwen3-TTS-12Hz-1.7B-Base

Manifest format (JSON):
    [
        {"speaker_id": "john", "text": "Hello world", "language": "English"},
        {"speaker_id": "jane", "text": "Testing batch", "language": "Auto"}
    ]
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import time
from typing import Any

import soundfile as sf
import torch

# vLLM-Omni requires spawn for multi-process workers.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

from embedding_cache import EmbeddingCacheManager

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Prompt length estimation (mirrors the pattern from end2end.py)
# ---------------------------------------------------------------------------

_ESTIMATOR_CACHE: dict[str, Any] = {}


def _estimate_prompt_len(additional_information: dict[str, Any], model_name: str) -> int:
    """Estimate placeholder prompt_token_ids length for the AR Talker.

    The Talker replaces all input embeddings via preprocess, so placeholder
    values are irrelevant -- but length must match the embeddings that
    preprocess will produce.
    """
    try:
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
            Qwen3TTSConfig,
        )
        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
            Qwen3TTSTalkerForConditionalGeneration,
        )

        if model_name not in _ESTIMATOR_CACHE:
            from transformers import AutoTokenizer

            tok = AutoTokenizer.from_pretrained(
                model_name, trust_remote_code=True, padding_side="left"
            )
            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
            _ESTIMATOR_CACHE[model_name] = (tok, getattr(cfg, "talker_config", None))

        tok, tcfg = _ESTIMATOR_CACHE[model_name]
        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
            additional_information=additional_information,
            task_type=task_type,
            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
            codec_language_id=getattr(tcfg, "codec_language_id", None),
            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
        )
    except Exception as exc:
        logger.warning("Prompt length estimation failed, using fallback 2048: %s", exc)
        return 2048


# ---------------------------------------------------------------------------
# Audio extraction helper
# ---------------------------------------------------------------------------


def extract_audio_from_output(output: Any) -> tuple:
    """Extract (audio_numpy, sample_rate) from an OmniRequestOutput."""
    mm = output.outputs[0].multimodal_output
    audio_key = "audio" if "audio" in mm else "model_outputs"
    audio_tensor = mm[audio_key]
    sr_val = mm.get("sr", 24000)
    if hasattr(sr_val, "item"):
        sr = sr_val.item()
    elif isinstance(sr_val, list):
        sr = int(sr_val[-1])
    else:
        sr = int(sr_val)

    # Async chunk mode returns a list of tensors.
    if isinstance(audio_tensor, list):
        audio_tensor = torch.cat(audio_tensor, dim=-1)

    audio_np = audio_tensor.float().detach().cpu().numpy()
    if audio_np.ndim > 1:
        audio_np = audio_np.flatten()

    return audio_np, sr


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def build_prompt(
    text: str,
    embedding: torch.Tensor,
    language: str,
    model_name: str,
    max_new_tokens: int,
) -> dict[str, Any]:
    """Build a single Omni prompt dict with a pre-computed embedding."""
    additional_information = {
        "task_type": ["Base"],
        "text": [text],
        "language": [language],
        "x_vector_only_mode": [True],
        "non_streaming_mode": [True],  # All text in prefill → fewer decode ops
        # Top-level tensor: survives additional_information serialization.
        "ref_spk_embedding": embedding,
        "max_new_tokens": [max_new_tokens],
    }
    ph_len = _estimate_prompt_len(additional_information, model_name)
    return {
        "prompt_token_ids": [0] * ph_len,
        "additional_information": additional_information,
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Batch TTS with pre-computed speaker embeddings"
    )
    parser.add_argument(
        "--manifest",
        required=True,
        help="Path to JSON manifest file with text/speaker pairs",
    )
    parser.add_argument(
        "--profiles-dir",
        default="voice_profiles",
        help="Directory containing .safetensors speaker profiles",
    )
    parser.add_argument(
        "--output-dir",
        default="output_audio",
        help="Directory to save output WAV files",
    )
    parser.add_argument(
        "--model",
        default="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
        help="HuggingFace model name or local path",
    )
    parser.add_argument(
        "--stage-configs-path",
        default="stage_config_optimized.yaml",
        help="Path to vLLM-Omni stage configuration YAML",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=2048,
        help="Maximum codec tokens to generate per utterance",
    )
    parser.add_argument(
        "--stage-init-timeout",
        type=int,
        default=300,
        help="Timeout in seconds for stage initialization",
    )
    parser.add_argument(
        "--log-stats",
        action="store_true",
        help="Enable vLLM-Omni statistics logging",
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    # ── 1. Load manifest ──────────────────────────────────────────────
    with open(args.manifest) as f:
        manifest = json.load(f)

    if not isinstance(manifest, list) or not manifest:
        logger.error("Manifest must be a non-empty JSON array")
        sys.exit(1)

    logger.info("Loaded %d items from manifest", len(manifest))

    # ── 2. Load embedding cache ───────────────────────────────────────
    cache = EmbeddingCacheManager(
        profiles_dir=args.profiles_dir,
        device="cpu",  # Keep on CPU; vLLM serializes to CPU anyway.
    )
    n_loaded = cache.preload_all()
    if n_loaded == 0:
        logger.error("No embeddings found in %s", args.profiles_dir)
        sys.exit(1)

    # Validate all speakers exist before starting expensive model load.
    for i, item in enumerate(manifest):
        sid = item.get("speaker_id")
        if sid not in cache:
            logger.error("Manifest item %d: speaker_id '%s' not in cache", i, sid)
            sys.exit(1)

    # ── 3. Build all prompts ──────────────────────────────────────────
    logger.info("Building prompts...")
    prompts = []
    for item in manifest:
        embedding = cache.get(item["speaker_id"])
        text = item["text"]
        language = item.get("language", "Auto")
        prompt = build_prompt(
            text=text,
            embedding=embedding,
            language=language,
            model_name=args.model,
            max_new_tokens=args.max_new_tokens,
        )
        prompts.append(prompt)

    logger.info("Built %d prompts", len(prompts))

    # ── 4. Initialize Omni engine ─────────────────────────────────────
    from vllm_omni import Omni

    logger.info("Initializing Omni engine with model=%s", args.model)
    omni = Omni(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )

    # ── 5. Generate ───────────────────────────────────────────────────
    os.makedirs(args.output_dir, exist_ok=True)

    logger.info("Starting batch generation...")
    gen_start = time.perf_counter()
    results: dict[str, tuple] = {}  # request_id -> (audio_np, sr)

    omni_generator = omni.generate(prompts, sampling_params_list=None)
    for stage_outputs in omni_generator:
        for output in stage_outputs.request_output:
            request_id = output.request_id
            try:
                audio_np, sr = extract_audio_from_output(output)
                results[request_id] = (audio_np, sr)
            except Exception as e:
                logger.error("Failed to extract audio for %s: %s", request_id, e)

    gen_elapsed = time.perf_counter() - gen_start

    # ── 6. Save WAV files ─────────────────────────────────────────────
    saved = 0
    for i, item in enumerate(manifest):
        # Omni assigns sequential request IDs; match by index.
        rid = str(i)
        if rid not in results:
            # Try other common ID formats
            for candidate in results:
                if candidate.endswith(f"_{i}") or candidate == str(i):
                    rid = candidate
                    break

        if rid in results:
            audio_np, sr = results[rid]
            output_name = item.get("output", f"output_{i}.wav")
            output_path = os.path.join(args.output_dir, output_name)
            sf.write(output_path, audio_np, samplerate=sr, format="WAV")
            logger.info("Saved: %s (%.1fs audio)", output_path, len(audio_np) / sr)
            saved += 1
        else:
            logger.warning("No output for manifest item %d (speaker=%s)", i, item.get("speaker_id"))

    # ── 7. Print statistics ───────────────────────────────────────────
    total_audio_s = sum(
        len(a) / s for a, s in results.values()
    )
    logger.info(
        "Batch complete: %d/%d saved, %.1fs generation time, "
        "%.1fs total audio, %.2fx realtime",
        saved,
        len(manifest),
        gen_elapsed,
        total_audio_s,
        total_audio_s / gen_elapsed if gen_elapsed > 0 else 0,
    )
    logger.info("Cache stats: %s", cache.stats())


if __name__ == "__main__":
    main()