vllm-omni/tools/breeze_tts_2_upstream_golden.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Generate deterministic Breeze-TTS-2 golden frames with the upstream runtime.

This tool is intentionally separate from vLLM-Omni: it imports the unmodified
``breezeblue-ai/breeze-tts`` reference implementation, fixes greedy decoding,
and writes both generated codec frames and decoded audio for offline parity
checks. It is a diagnostic tool, not a serving path.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from types import SimpleNamespace

import soundfile as sf
import torch


def _plain_template(upstream_templates):
    """The upstream CLI does not register tts_plain; reuse its segment builder."""
    return SimpleNamespace(
        name="tts_plain",
        required_fields=("text",),
        build_segments=upstream_templates._tts_plain_segments,
        build_negative_segments=None,
        build_dual_branches=None,
    )


def _greedy(model) -> None:
    model.generation_config.do_sample = False
    model.generation_config.temperature = None
    model.generation_config.top_p = None
    model.generation_config.top_k = None
    depth = model.depth_decoder.generation_config
    depth.depth_decoder_do_sample = False
    depth.do_sample = False
    depth.temperature = None
    depth.top_p = None
    depth.top_k = None


def _valid_frames(sequences: torch.Tensor, pad_token_id: int) -> torch.Tensor:
    frames = sequences[0].detach().cpu().to(torch.long)
    if frames.ndim == 1:
        frames = frames.unsqueeze(0)
    is_pad = (frames == pad_token_id).all(dim=-1)
    if bool(is_pad.any()):
        frames = frames[: int(torch.argmax(is_pad.int()).item())]
    return frames.contiguous()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--upstream", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--reference-audio", type=Path)
    parser.add_argument("--max-new-tokens", type=int, default=32)
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)
    upstream_posix = str(args.upstream.resolve())
    if upstream_posix not in sys.path:
        sys.path.insert(0, upstream_posix)

    from breeze_infer.runtime import load_runtime, resolve_device
    from breeze_infer.templates import get_template, prepare_inputs

    device = resolve_device()
    # ``torch.manual_seed`` seeds CUDA RNGs as well, so no device-specific
    # seeding call is needed here.
    torch.manual_seed(0)
    tokenizer, model, audio_tokenizer = load_runtime(
        args.model,
        device=device,
        attn_implementation="eager",
    )
    # Do not enable Transformers 4.x's ``fix_mistral_regex`` compatibility
    # path: vLLM-Omni runs with Transformers 5.x, and direct tokenization checks
    # showed both runtimes produce the same token IDs without that patch.
    _greedy(model)
    model.eval()
    first_backbone_hidden: list[torch.Tensor] = []

    def _capture_first_backbone_hidden(_module, _inputs, output):
        if first_backbone_hidden:
            return
        hidden = getattr(output, "last_hidden_state", None)
        if isinstance(hidden, torch.Tensor):
            first_backbone_hidden.append(hidden[0, -1].detach().cpu().to(torch.float32).contiguous())

    model.backbone_model.register_forward_hook(_capture_first_backbone_hidden)

    cases = [
        {
            "name": "plain",
            "template": _plain_template(sys.modules["breeze_infer.templates"]),
            "request": {
                "id": "golden-plain",
                "text": "Golden plain alignment check.",
                "speaker": "S0",
            },
        },
        {
            "name": "instruction",
            "template": get_template("tts_instruction"),
            "request": {
                "id": "golden-instruction",
                "text": "Golden instruction alignment check.",
                "instruction": "Speak calmly, warmly, and clearly.",
                "speaker": "S0",
            },
        },
    ]
    if args.reference_audio is not None:
        cases.append(
            {
                "name": "reference_edit",
                "template": get_template("ref_edit_tata"),
                "request": {
                    "id": "golden-reference-edit",
                    "text": "Golden reference alignment check.",
                    "instruction": "Speak calmly and clearly.",
                    "ref_audio_path": str(args.reference_audio),
                    "ref_text": "Hello from Breeze.",
                    "speaker": "S0",
                },
            }
        )

    manifest = {
        "runtime": "upstream breezeblue-ai/breeze-tts",
        "model": str(args.model.resolve()),
        "upstream": upstream_posix,
        "device": device,
        "seed": 0,
        "max_new_tokens": args.max_new_tokens,
        "decoding": "greedy",
        "cfg_scale": 1.0,
        "cases": [],
    }
    for case in cases:
        first_backbone_hidden.clear()
        inputs = prepare_inputs(
            tokenizer,
            audio_tokenizer,
            model,
            [case["request"]],
            case["template"],
            guidance_scale=1.0,
            guidance_scale_ref=None,
            guidance_scale_ins=None,
        )
        output = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,
            return_dict_in_generate=True,
            output_audio=True,
            audio_tokenizer=audio_tokenizer,
        )
        frames = _valid_frames(output.sequences, model.config.codebook_pad_token_id)
        audio = output.audio[0].detach().cpu().to(torch.float32).reshape(-1)
        case_dir = args.output_dir / case["name"]
        case_dir.mkdir(parents=True, exist_ok=True)
        torch.save(frames, case_dir / "codes.pt")
        prompt_ids = inputs["input_ids"].detach().cpu().to(torch.long).reshape(-1)
        torch.save(prompt_ids, case_dir / "prompt_ids.pt")
        if first_backbone_hidden:
            torch.save(first_backbone_hidden[0], case_dir / "prefill_hidden.pt")
        sf.write(case_dir / "upstream.wav", audio.numpy(), 24_000, subtype="PCM_16")
        entry = {
            "name": case["name"],
            "request": case["request"],
            "prompt_tokens": int(prompt_ids.numel()),
            "frames": int(frames.shape[0]),
            "codebooks": int(frames.shape[-1]) if frames.ndim == 2 else 0,
            "min_code": int(frames.min().item()) if frames.numel() else None,
            "max_code": int(frames.max().item()) if frames.numel() else None,
            "audio_samples": int(audio.numel()),
            "audio_sha256": __import__("hashlib").sha256((case_dir / "upstream.wav").read_bytes()).hexdigest(),
        }
        manifest["cases"].append(entry)
        print(json.dumps(entry), flush=True)

    (args.output_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()