inference.py

#!/usr/bin/env python3
"""
Generate audio for each input text with the Priyanka LoRA and full-finetune checkpoints.

Examples:
    python inference.py --text "नमस्ते, आप कैसे हैं?" --text "This is a second sample."

    python inference.py --text-file texts.txt --output-dir outputs/priyanka_inference

    python inference.py \
        --text "This is voice cloning." \
        --prompt-audio examples/reference_speaker.wav \
        --prompt-text "Reference speaker transcript."
"""

from __future__ import annotations

import argparse
import gc
import json
import re
import sys
from dataclasses import dataclass
from pathlib import Path

import soundfile as sf

ROOT = Path(__file__).resolve().parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from voxcpm.core import VoxCPM  # noqa: E402
from voxcpm.model.voxcpm import LoRAConfig  # noqa: E402


CHECKPOINTS_ROOT = ROOT / "checkpoints"
DEFAULT_CHECKPOINTS = (
    {
        "name": "03_priyanka_lora_step_0000999",
        "kind": "lora",
        "path": CHECKPOINTS_ROOT / "03_priyanka_lora" / "step_0000999",
    },
    {
        "name": "04_priyanka_full_step_0000500",
        "kind": "full",
        "path": CHECKPOINTS_ROOT / "04_priyanka_full" / "step_0000500",
    },
)


@dataclass(frozen=True)
class CheckpointSpec:
    name: str
    kind: str
    path: Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate WAV files for every provided text using two Priyanka VoxCPM checkpoints."
    )
    parser.add_argument(
        "--text",
        action="append",
        default=[],
        help="Text to synthesize. Can be passed multiple times.",
    )
    parser.add_argument(
        "--text-file",
        type=Path,
        default=None,
        help="UTF-8 text file with one synthesis text per non-empty line.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=ROOT / "outputs" / "priyanka_inference",
        help="Directory where generated WAV files and manifest.json are saved.",
    )
    parser.add_argument(
        "--prompt-audio",
        type=Path,
        default=None,
        help="Optional reference/prompt WAV path. Must be used with --prompt-text.",
    )
    parser.add_argument(
        "--prompt-text",
        default=None,
        help="Transcript for --prompt-audio.",
    )
    parser.add_argument(
        "--reference-audio",
        type=Path,
        default=None,
        help="Optional VoxCPM2 reference WAV path for voice cloning.",
    )
    parser.add_argument("--cfg-value", type=float, default=2.0, help="CFG scale.")
    parser.add_argument(
        "--inference-timesteps",
        type=int,
        default=10,
        help="Number of diffusion inference steps.",
    )
    parser.add_argument("--max-len", type=int, default=600, help="Maximum generation length.")
    parser.add_argument("--normalize", action="store_true", help="Enable text normalization.")
    parser.add_argument(
        "--device",
        default=None,
        help="Runtime device, for example auto, cuda, cuda:0, mps, or cpu.",
    )
    parser.add_argument(
        "--no-optimize",
        action="store_true",
        help="Disable VoxCPM warmup/optimization.",
    )
    parser.add_argument(
        "--denoise",
        action="store_true",
        help="Run denoiser for prompt/reference audio. The denoiser is loaded only when this is set.",
    )
    return parser.parse_args()


def read_texts(args: argparse.Namespace) -> list[str]:
    texts = [text.strip() for text in args.text if text and text.strip()]

    if args.text_file:
        if not args.text_file.exists():
            raise FileNotFoundError(f"--text-file does not exist: {args.text_file}")
        file_texts = [
            line.strip()
            for line in args.text_file.read_text(encoding="utf-8").splitlines()
            if line.strip()
        ]
        texts.extend(file_texts)

    if not texts:
        raise ValueError("Provide at least one input with --text or --text-file.")

    return texts


def validate_audio_args(args: argparse.Namespace) -> None:
    if (args.prompt_audio is None) != (args.prompt_text is None):
        raise ValueError("--prompt-audio and --prompt-text must be provided together.")
    for path_arg in (args.prompt_audio, args.reference_audio):
        if path_arg is not None and not path_arg.exists():
            raise FileNotFoundError(f"Audio file does not exist: {path_arg}")


def load_lora_model(ckpt_dir: Path, args: argparse.Namespace) -> VoxCPM:
    lora_config_path = ckpt_dir / "lora_config.json"
    if not lora_config_path.exists():
        raise FileNotFoundError(f"Missing LoRA config: {lora_config_path}")

    with lora_config_path.open("r", encoding="utf-8") as f:
        lora_info = json.load(f)

    base_model = lora_info.get("base_model")
    if not base_model:
        raise ValueError(f"'base_model' is missing in {lora_config_path}")

    lora_cfg = LoRAConfig(**lora_info.get("lora_config", {}))
    return VoxCPM.from_pretrained(
        hf_model_id=base_model,
        load_denoiser=args.denoise,
        optimize=not args.no_optimize,
        device=args.device,
        lora_config=lora_cfg,
        lora_weights_path=str(ckpt_dir),
    )


def load_full_model(ckpt_dir: Path, args: argparse.Namespace) -> VoxCPM:
    return VoxCPM.from_pretrained(
        hf_model_id=str(ckpt_dir),
        load_denoiser=args.denoise,
        optimize=not args.no_optimize,
        device=args.device,
    )


def load_model(spec: CheckpointSpec, args: argparse.Namespace) -> VoxCPM:
    if not spec.path.exists():
        raise FileNotFoundError(f"Checkpoint does not exist: {spec.path}")
    if spec.kind == "lora":
        return load_lora_model(spec.path, args)
    if spec.kind == "full":
        return load_full_model(spec.path, args)
    raise ValueError(f"Unsupported checkpoint kind: {spec.kind}")


def slugify(value: str, max_len: int = 48) -> str:
    value = re.sub(r"\s+", "_", value.strip().lower())
    value = re.sub(r"[^0-9a-zA-Z_\-]+", "", value)
    value = value.strip("_-")
    return (value[:max_len].strip("_-") or "text")


def output_path(output_dir: Path, spec: CheckpointSpec, index: int, text: str) -> Path:
    text_slug = slugify(text)
    return output_dir / f"text_{index:03d}_{spec.name}_{text_slug}.wav"


def release_model(model: VoxCPM) -> None:
    del model
    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def main() -> None:
    args = parse_args()
    texts = read_texts(args)
    validate_audio_args(args)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    specs = [CheckpointSpec(**item) for item in DEFAULT_CHECKPOINTS]
    manifest = {
        "output_dir": str(args.output_dir),
        "texts": texts,
        "checkpoints": [],
        "files": [],
    }

    print(f"Generating {len(texts)} text(s) with {len(specs)} checkpoint(s).", file=sys.stderr)

    for spec in specs:
        print(f"\nLoading {spec.name}: {spec.path}", file=sys.stderr)
        model = load_model(spec, args)
        sample_rate = model.tts_model.sample_rate
        manifest["checkpoints"].append(
            {"name": spec.name, "kind": spec.kind, "path": str(spec.path), "sample_rate": sample_rate}
        )

        try:
            for index, text in enumerate(texts, start=1):
                wav_path = output_path(args.output_dir, spec, index, text)
                print(f"[{spec.name}] text {index}/{len(texts)} -> {wav_path}", file=sys.stderr)

                audio_np = model.generate(
                    text=text,
                    prompt_wav_path=str(args.prompt_audio) if args.prompt_audio else None,
                    prompt_text=args.prompt_text,
                    reference_wav_path=str(args.reference_audio) if args.reference_audio else None,
                    cfg_value=args.cfg_value,
                    inference_timesteps=args.inference_timesteps,
                    max_len=args.max_len,
                    normalize=args.normalize,
                    denoise=args.denoise,
                )
                sf.write(str(wav_path), audio_np, sample_rate)
                manifest["files"].append(
                    {
                        "checkpoint": spec.name,
                        "text_index": index,
                        "text": text,
                        "path": str(wav_path),
                        "duration_seconds": len(audio_np) / sample_rate,
                    }
                )
        finally:
            release_model(model)

    manifest_path = args.output_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"\nDone. Wrote {len(manifest['files'])} WAV file(s).", file=sys.stderr)
    print(f"Manifest: {manifest_path}", file=sys.stderr)


if __name__ == "__main__":
    main()