# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import os
import urllib.request
from pathlib import Path

import numpy as np
import soundfile as sf
from vllm import SamplingParams
from vllm.multimodal.media.audio import load_audio

from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config

# Upstream zero-shot reference clip
ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"


def _default_ref_audio() -> str:
    # Download the upstream zero_shot_prompt.wav into the current dir
    dest = Path("zero_shot_prompt.wav")
    if not dest.exists() or dest.stat().st_size == 0:
        print(f"Downloading default reference audio to {dest}")
        urllib.request.urlretrieve(ZERO_SHOT_PROMPT_URL, dest)

    return str(dest)


def run_e2e():
    parser = argparse.ArgumentParser()
    # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help="Override the deploy config path. If unset, auto-loads "
        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
    )
    parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
    parser.add_argument(
        "--prompt-text",
        type=str,
        default="You are a helpful assistant.<|endofprompt|>希望你以后，能够做的比我还好呦!",
    )
    parser.add_argument(
        "--ref-audio",
        type=str,
        default=None,
        help="Path to reference audio for voice cloning. "
        "If unset, downloads the upstream CosyVoice3 zero-shot prompt audio clip",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        required=True,
        help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
    )
    args = parser.parse_args()
    # Ensure tokenizer directory exists
    if not os.path.exists(args.tokenizer):
        raise FileNotFoundError(f"{args.tokenizer} does not exist!")

    if args.deploy_config is not None and not os.path.exists(args.deploy_config):
        raise FileNotFoundError(f"{args.deploy_config} does not exist!")

    print(f"Initializing cosyvoice E2E with model={args.model}")

    omni = Omni(
        model=args.model,
        deploy_config=args.deploy_config,
        tokenizer=args.tokenizer,
        log_stats=True,
    )

    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}

    print("Model initialized. Preparing inputs...")
    ref_audio_path = args.ref_audio or _default_ref_audio()
    if not os.path.exists(ref_audio_path):
        raise FileNotFoundError(f"Audio file not found: {ref_audio_path}")
    # Load at native sample rate
    audio_signal, sr = load_audio(ref_audio_path, sr=None)

    # Validate sample rate before processing (similar to original CosyVoice)
    min_sr = 16000
    if sr < min_sr:
        raise ValueError(
            f"Audio sample rate {sr} Hz is too low. "
            f"Minimum required: {min_sr} Hz. "
            f"Please provide audio with sample rate >= {min_sr} Hz."
        )

    audio_data = (audio_signal.astype(np.float32), sr)

    prompts = {
        "prompt": args.text,
        "multi_modal_data": {
            "audio": audio_data,
        },
        "mm_processor_kwargs": {
            "prompt_text": args.prompt_text,
            "sample_rate": audio_data[1],
        },
    }

    print(f"Generating for prompt: {args.text}")

    config = CosyVoice3Config()
    tokenizer = get_qwen_tokenizer(
        token_path=args.tokenizer,
        skip_special_tokens=config.skip_special_tokens,
        version=config.version,
    )
    _, text_token_len = extract_text_token(args.text, tokenizer, config.allowed_special)
    base_len = int(text_token_len)
    min_len = int(base_len * config.min_token_text_ratio)
    max_len = int(base_len * config.max_token_text_ratio)

    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
    gpt_sampling = SamplingParams(
        temperature=1.0,
        top_p=sampling_cfg["top_p"],
        top_k=sampling_cfg["top_k"],
        repetition_penalty=2.0,
        min_tokens=min_len,
        max_tokens=max_len,
        stop_token_ids=[sampling_cfg["eos_token_id"]],
        # allowed_token_ids=list(range(6561+3)),
        detokenize=False,
    )
    # Not used
    s2mel_sampling = SamplingParams(
        temperature=1.0,
        top_p=1.0,
        top_k=-1,
        repetition_penalty=2.0,
        max_tokens=256,
        detokenize=False,
    )

    sampling_params_list = [gpt_sampling, s2mel_sampling]

    # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
        print("Starting profiler...")
        omni.start_profile()

    # Generate (Omni orchestrator requires a per-stage SamplingParams list)
    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))

    # Stop profiling and get results
    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
        print("Stopping profiler...")
        profile_results = omni.stop_profile()
        print(f"Profile traces saved to: {profile_results}")

    print(outputs)
    # Verify outputs
    print(f"Received {len(outputs)} outputs.")
    for i, output in enumerate(outputs):
        try:
            ro = output.request_output
            if ro is None:
                print("No request_output found.")
                continue

            # Multimodal output may be attached to RequestOutput or CompletionOutput.
            mm = getattr(ro, "multimodal_output", None)
            if not mm and ro.outputs:
                mm = getattr(ro.outputs[0], "multimodal_output", None)

            if mm:
                print(f"Multimodal output keys: {mm.keys()}")
                if "audio" in mm:
                    audio_out = mm["audio"]
                    print(f"Generated Audio Shape: {audio_out.shape}")
                    out_path = f"output_{i}.wav"
                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
                    print(f"Saved audio to {out_path}")
            else:
                print("No multimodal output found.")
        except Exception as e:
            print(f"Error inspecting output: {e}")
    omni.close()


if __name__ == "__main__":
    run_e2e()

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""GLM-TTS End-to-End Offline Inference Example.

GLM-TTS is a two-stage TTS system:
  - Stage 0 (AR): Llama-based model generates speech tokens from text
  - Stage 1 (DiT): Flow matching model converts speech tokens to audio

Usage:
    # Sync two-stage (default)
    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
        --model /path/to/GLM-TTS \
        --text "你好，这是一个语音合成测试。" \
        --ref-audio /path/to/reference.wav \
        --ref-text "参考音频的转录文本。" \
        --output-dir ./output

    # Async chunk mode (streaming DiT)
    python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
        --model /path/to/GLM-TTS --async-chunk \
        --text "你好，这是一个语音合成测试。" \
        --ref-audio /path/to/reference.wav \
        --ref-text "参考音频的转录文本。" \
        --output-dir ./output
"""

import base64
import io
import logging
import os
import tempfile
import time
from typing import Any
from urllib.request import urlopen

import numpy as np
import soundfile as sf
import torch
import yaml

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

from vllm.utils.argparse_utils import FlexibleArgumentParser

from vllm_omni import Omni
from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata

logger = logging.getLogger(__name__)

DEFAULT_DEPLOY_CONFIG = os.path.join(
    os.path.dirname(__file__),
    "..",
    "..",
    "..",
    "..",
    "vllm_omni",
    "deploy",
    "glm_tts.yaml",
)
SAMPLE_RATE = 24000


def _load_ref_audio(ref_audio: str) -> tuple[torch.Tensor, int]:
    """Load reference audio from file path, URL, or data URI."""
    if ref_audio.startswith(("http://", "https://")):
        with urlopen(ref_audio, timeout=60) as response:
            audio_obj: Any = io.BytesIO(response.read())
    elif ref_audio.startswith("data:"):
        _, _, encoded = ref_audio.partition(",")
        audio_obj = io.BytesIO(base64.b64decode(encoded))
    else:
        audio_obj = ref_audio
    wav_np, sr = sf.read(audio_obj, dtype="float32")
    if wav_np.ndim > 1:
        wav_np = wav_np.mean(axis=1)
    return torch.from_numpy(wav_np), int(sr)


def _concat_audio(audio_val: Any) -> np.ndarray:
    """Concatenate audio tensors from multimodal output."""
    if isinstance(audio_val, list):
        tensors = [torch.as_tensor(t).float().reshape(-1) for t in audio_val if t is not None]
        if not tensors:
            return np.zeros((0,), dtype=np.float32)
        return torch.cat(tensors, dim=-1).cpu().numpy().astype(np.float32, copy=False)
    if isinstance(audio_val, torch.Tensor):
        return audio_val.float().cpu().numpy().reshape(-1)
    return np.asarray(audio_val, dtype=np.float32).reshape(-1)


def _extract_sample_rate(audio_mm: dict) -> int:
    """Extract sample rate from multimodal output dict."""
    sr_raw = audio_mm.get("sr", SAMPLE_RATE)
    if isinstance(sr_raw, list):
        sr_raw = sr_raw[-1] if sr_raw else SAMPLE_RATE
    if hasattr(sr_raw, "item"):
        return int(sr_raw.item())
    return int(sr_raw)


def _modify_deploy_config(base_path: str, async_chunk: bool) -> str:
    """Build deploy config with explicit sync/async mode and eager execution.

    Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts_expansion.py``
    (``_get_deploy_config``) so that example runs match CI behavior.
    """
    with open(base_path) as f:
        cfg = yaml.safe_load(f)
    cfg["async_chunk"] = async_chunk
    for stage in cfg.get("stages", []):
        stage["enforce_eager"] = True
        if stage.get("stage_id") == 0:
            stage["async_scheduling"] = bool(async_chunk)
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False, prefix="glm_tts_")
    yaml.dump(cfg, tmp)
    tmp.close()
    return tmp.name


def main(args):
    """Run offline GLM-TTS inference."""
    os.makedirs(args.output_dir, exist_ok=True)
    base_deploy_config = args.deploy_config or DEFAULT_DEPLOY_CONFIG
    deploy_config_path = _modify_deploy_config(base_deploy_config, args.async_chunk)

    ref_audio_wav, ref_audio_sr = _load_ref_audio(args.ref_audio)
    if not args.ref_text:
        raise ValueError("GLM-TTS requires --ref-audio and --ref-text.")

    inputs = [
        {
            "prompt": args.text,
            "multi_modal_data": {
                "audio": (ref_audio_wav.float().cpu().numpy(), ref_audio_sr),
            },
            "modalities": ["audio"],
            "mm_processor_kwargs": {"prompt_text": args.ref_text},
            "additional_information": build_glm_tts_prefill_metadata(
                args.model,
                args.text,
                args.ref_text,
            ),
        }
    ]

    omni = Omni(
        model=args.model,
        stage_configs_path=deploy_config_path,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )

    t_start = time.perf_counter()
    outputs = omni.generate(inputs)
    elapsed = (time.perf_counter() - t_start) * 1000

    assert outputs, "No outputs returned"
    audio_mm = outputs[0].multimodal_output
    assert "audio" in audio_mm, "No audio output found"

    audio = _concat_audio(audio_mm["audio"])
    sr = _extract_sample_rate(audio_mm)
    out_path = os.path.join(args.output_dir, "output.wav")
    sf.write(out_path, audio, samplerate=sr, format="WAV")

    logger.info("Saved %s (%.2fs @ %dHz)", out_path, len(audio) / sr, sr)
    logger.info("Total inference: %.1f ms", elapsed)


def parse_args():
    parser = FlexibleArgumentParser(description="GLM-TTS Text-to-Speech Example")
    parser.add_argument("--model", type=str, required=True, help="Model path")
    parser.add_argument("--text", type=str, default="你好，这是一个语音合成测试。")
    parser.add_argument("--output-dir", type=str, default="./output")
    parser.add_argument("--ref-audio", type=str, required=True, help="Reference WAV path/URL")
    parser.add_argument("--ref-text", type=str, required=True, help="Transcript of ref audio")
    parser.add_argument("--deploy-config", type=str, default=None)
    parser.add_argument(
        "--async-chunk",
        action="store_true",
        default=False,
        help="Enable async_chunk mode (streaming DiT). Default: sync two-stage.",
    )
    parser.add_argument("--log-stats", action="store_true")
    parser.add_argument("--stage-init-timeout", type=int, default=600)
    return parser.parse_args()


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    main(parse_args())

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline higgs-audio v2 inference example.

Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
through the vLLM-Omni engine without going through the HTTP server, and
saves a 24 kHz mono WAV per prompt.

Example:

    python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
        --texts "Hello world." \\
                "The quick brown fox jumps over the lazy dog." \\
        --output-dir results/wavs
"""

from __future__ import annotations

import os

# DeepGEMM FP8 kernels require an optional backend that may not be installed.
# Disable the warmup before importing vLLM so engine startup falls back to the
# regular gemm path. Users with deep_gemm installed can override these.
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")

import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser

SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
    "Hello world.",
    "The quick brown fox jumps over the lazy dog.",
)


def parse_args():
    parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
    parser.add_argument(
        "--model",
        type=str,
        default="bosonai/higgs-audio-v2-generation-3B-base",
        help="Stage-0 talker model id or path.",
    )
    parser.add_argument(
        "--texts",
        type=str,
        nargs="+",
        default=list(DEFAULT_TEXTS),
        help="One or more plain-text prompts to synthesize.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="results/wavs",
        help="Directory to write per-prompt WAV files.",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help="Override the deploy config path. Auto-loads "
        "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=500,
        help="Cap on Stage-0 codec frames per request.",
    )
    parser.add_argument(
        "--ref-audio",
        type=str,
        default=None,
        help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
    )
    parser.add_argument(
        "--ref-text",
        type=str,
        default=None,
        help="Transcript of the reference clip. Required when --ref-audio is set.",
    )
    return parser.parse_args()


def _slugify(text: str) -> str:
    import re

    slug = re.sub(r"\s+", "_", text.strip().lower())
    slug = re.sub(r"[^a-z0-9_]+", "", slug)
    return slug[:48] or "prompt"


def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
    """Pull the final concatenated PCM tensor out of a request's multimodal_output."""
    audio = multimodal_output.get("model_outputs")
    if audio is None:
        audio = multimodal_output.get("audio")
    if audio is None:
        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
    if isinstance(audio, list):
        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
        if not valid:
            raise ValueError("audio list is empty")
        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
    return torch.as_tensor(audio).float().cpu().reshape(-1)


def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
    arr = pcm.numpy()
    if arr.dtype.kind == "f":
        arr = np.clip(arr, -1.0, 1.0)
        arr = (arr * 32767.0).astype(np.int16)
    else:
        arr = arr.astype(np.int16)
    return arr


def main():
    args = parse_args()
    if (args.ref_audio is None) != (args.ref_text is None):
        raise SystemExit("--ref-audio and --ref-text must be supplied together")
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    engine = Omni(model=args.model, deploy_config=args.deploy_config)

    # Build prompt_token_ids using the same path serving_speech.py takes online.
    from transformers import AutoProcessor

    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
        build_plain_text_prompt,
        build_voice_clone_prompt,
        input_ids_to_python_list,
    )

    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)

    # Voice-clone path: load the reference clip once. The HF processor will
    # encode it via the bundled HiggsAudioV2TokenizerModel each time we build
    # a prompt below. This is cheap on CPU for a few-second clip.
    ref_wav: np.ndarray | None = None
    ref_sr: int | None = None
    if args.ref_audio is not None:
        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
        if ref_wav.ndim == 2:
            ref_wav = ref_wav.mean(axis=1)

    print(f"Model       : {args.model}")
    print(f"Prompts     : {len(args.texts)}")
    print(f"Output dir  : {output_dir}")
    print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")

    # Run one prompt at a time. The Stage-0 talker's per-slot audio state is
    # request-scoped; submitting multiple prompts in the same engine.generate()
    # call would batch them in the AR runner and exercise a code path that is
    # not validated for this model yet.
    total_elapsed = 0.0
    total_dur = 0.0
    for text in args.texts:
        if ref_wav is not None:
            out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
            prompt = {
                "prompt_token_ids": out["prompt_token_ids"],
                # Bare tensors (NOT list-wrapped): the msgspec serializer in
                # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
                # and list[Tensor] → list_data (which silently strips tensors).
                "additional_information": {
                    "audio_input_ids": out["audio_input_ids"],
                    "audio_input_ids_mask": out["audio_input_ids_mask"],
                },
            }
        else:
            inputs = build_plain_text_prompt(processor, text)
            prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
        t_start = time.perf_counter()
        outputs = engine.generate([prompt])
        elapsed = time.perf_counter() - t_start
        total_elapsed += elapsed

        mm = outputs[0].outputs[0].multimodal_output
        pcm = _extract_pcm(mm)
        slug = _slugify(text)
        out_path = output_dir / f"{slug}.wav"
        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
        dur = pcm.numel() / SAMPLE_RATE
        total_dur += dur
        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")

    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")


if __name__ == "__main__":
    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
    main()

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline higgs-audio v3 inference example.

Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
through the vLLM-Omni engine without going through the HTTP server, and
saves a 24 kHz mono WAV per prompt.

Example (plain TTS):

    python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
        --texts "Hello world." \\
        --output-dir results/higgs_v3_wavs

Example (voice clone):

    python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
        --texts "Hello world." \\
        --ref-audio path/to/reference.wav \\
        --ref-text "Transcript of the reference clip." \\
        --output-dir results/higgs_v3_wavs
"""

from __future__ import annotations

import os

os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")

import re
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from vllm_omni import Omni

SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
    "Hello world.",
    "The quick brown fox jumps over the lazy dog.",
    "Today is a beautiful day for a walk in the park.",
)


def parse_args():
    import argparse

    parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
    parser.add_argument(
        "--model",
        type=str,
        default="bosonai/higgs-audio-v3-tts-4b",
        help="Stage-0 talker model id or path.",
    )
    parser.add_argument(
        "--texts",
        type=str,
        nargs="+",
        default=list(DEFAULT_TEXTS),
        help="One or more plain-text prompts to synthesize.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="results/higgs_v3_wavs",
        help="Directory to write per-prompt WAV files.",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help="Override the deploy config path.",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=2048,
        help="Cap on Stage-0 codec frames per request.",
    )
    parser.add_argument(
        "--ref-audio",
        type=str,
        default=None,
        help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
    )
    parser.add_argument(
        "--ref-text",
        type=str,
        default=None,
        help="Transcript of the reference audio. Optional but improves fidelity.",
    )
    return parser.parse_args()


def _slugify(text: str) -> str:
    slug = re.sub(r"\s+", "_", text.strip().lower())
    slug = re.sub(r"[^a-z0-9_]+", "", slug)
    return slug[:48] or "prompt"


def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
    audio = multimodal_output.get("model_outputs")
    if audio is None:
        audio = multimodal_output.get("audio")
    if audio is None:
        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
    if isinstance(audio, list):
        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
        if not valid:
            raise ValueError("audio list is empty")
        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
    return torch.as_tensor(audio).float().cpu().reshape(-1)


def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
    arr = pcm.numpy()
    if arr.dtype.kind == "f":
        arr = np.clip(arr, -1.0, 1.0)
        arr = (arr * 32767.0).astype(np.int16)
    else:
        arr = arr.astype(np.int16)
    return arr


def main():
    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)

    from transformers import AutoTokenizer

    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
        HiggsAudioV3TokenizerAdapter,
        apply_delay_pattern,
        encode_reference_audio,
    )

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    adapter = HiggsAudioV3TokenizerAdapter(tokenizer)

    # Load and encode reference audio once (if voice cloning)
    ref_codes_delayed: torch.Tensor | None = None
    if args.ref_audio is not None:
        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
        if ref_wav.ndim == 2:
            ref_wav = ref_wav.mean(axis=1)
        ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
        ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
        print(f"Reference   : {args.ref_audio}")
        print(f"Ref codes   : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
        if args.ref_text:
            print(f"Ref text    : {args.ref_text}")

    print(f"Model       : {args.model}")
    print(f"Prompts     : {len(args.texts)}")
    print(f"Output dir  : {output_dir}")
    print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")

    total_elapsed = 0.0
    total_dur = 0.0
    for text in args.texts:
        if ref_codes_delayed is not None:
            prompt_ids = adapter.build_prompt(
                text,
                num_ref_tokens=int(ref_codes_delayed.shape[0]),
                reference_text=args.ref_text,
            )
            prompt = {
                "prompt_token_ids": prompt_ids,
                "additional_information": {
                    "audio_input_ids": ref_codes_delayed.to(torch.long),
                    "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
                },
            }
        else:
            prompt_ids = adapter.build_prompt(text)
            prompt = {"prompt_token_ids": prompt_ids}

        t_start = time.perf_counter()
        outputs = engine.generate([prompt])
        elapsed = time.perf_counter() - t_start
        total_elapsed += elapsed

        mm = outputs[0].outputs[0].multimodal_output
        pcm = _extract_pcm(mm)
        slug = _slugify(text)
        out_path = output_dir / f"{slug}.wav"
        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
        dur = pcm.numel() / SAMPLE_RATE
        total_dur += dur
        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")

    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")


if __name__ == "__main__":
    main()

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline inference example for IndexTTS2 via vLLM-Omni.

Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
Output is 22050 Hz mono WAV.

Usage:
  python end2end.py \
    --model /path/to/IndexTeam/IndexTTS-2 \
    --text "你好，这是一个语音合成测试。" \
    --ref-audio /path/to/ref.wav

  # With emotion audio:
  python end2end.py \
    --model /path/to/IndexTeam/IndexTTS-2 \
    --text "今天天气真好！" \
    --ref-audio /path/to/ref.wav \
    --emo-audio /path/to/happy.wav
"""

from __future__ import annotations

import os
from pathlib import Path

import soundfile as sf
import torch
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser

os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")

from vllm_omni import Omni  # noqa: E402
from vllm_omni.model_executor.models.indextts2.prompt_utils import (  # noqa: E402
    build_indextts2_prefill_prompt_ids,
)


def build_request(
    model: str,
    text: str,
    ref_audio_path: str | None = None,
    emo_audio_path: str | None = None,
    emo_text: str | None = None,
    emo_vector: list[float] | None = None,
    emo_alpha: float | None = None,
    use_emo_text: bool = False,
    use_random: bool = False,
) -> dict:
    additional: dict = {"text": [text]}
    if ref_audio_path:
        additional["voice"] = [str(ref_audio_path)]
    if emo_audio_path:
        additional["emo_audio"] = [str(emo_audio_path)]
    if emo_text:
        additional["emo_text"] = [emo_text]
    if emo_vector is not None:
        additional["emo_vector"] = [emo_vector]
    if emo_alpha is not None:
        additional["emo_alpha"] = [emo_alpha]
    if use_emo_text:
        additional["use_emo_text"] = [True]
    if use_random:
        additional["use_random"] = [True]
    return {
        "prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
        "additional_information": additional,
    }


def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
    audio_np = waveform.float().numpy()
    sf.write(path, audio_np, sample_rate)
    print(f"  Saved {path} ({audio_np.shape}, {sample_rate} Hz)")


def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
    audio = mm.get("audio")
    if audio is None:
        audio = mm.get("model_outputs")
    if isinstance(audio, list):
        chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
        audio = torch.cat(chunks, dim=0) if chunks else None

    sr_val = mm.get("sr")
    if isinstance(sr_val, list):
        sr_val = sr_val[-1] if sr_val else None
    if hasattr(sr_val, "item"):
        sample_rate = int(sr_val.item())
    else:
        sample_rate = int(sr_val) if sr_val is not None else 22050

    return audio if isinstance(audio, torch.Tensor) else None, sample_rate


def main(args) -> None:
    omni = Omni(
        model=args.model,
        deploy_config=args.deploy_config,
        stage_init_timeout=args.stage_init_timeout,
    )

    # Stage 0: GPT AR (text → mel codes)
    gpt_sampling = SamplingParams(
        temperature=0.8,
        top_p=0.8,
        top_k=30,
        max_tokens=1500,
        repetition_penalty=10.0,
        stop_token_ids=[8193],
        seed=args.seed if args.seed is not None else 42,
        detokenize=False,
    )
    # Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
    s2mel_sampling = SamplingParams(
        temperature=0.0,
        max_tokens=65536,
        detokenize=True,
    )
    sampling_params = [gpt_sampling, s2mel_sampling]

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Synthesizing: {args.text!r}")
    if args.ref_audio:
        print(f"  ref_audio: {args.ref_audio}")
    inputs = build_request(
        model=args.model,
        text=args.text,
        ref_audio_path=args.ref_audio,
        emo_audio_path=args.emo_audio,
        emo_text=args.emo_text,
        emo_vector=args.emo_vector,
        emo_alpha=args.emo_alpha,
        use_emo_text=args.use_emo_text,
        use_random=args.use_random,
    )

    for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
        mm = omni_out.multimodal_output
        if not mm:
            print(f"  [req {i}] No multimodal output.")
            continue
        audio, sr = extract_audio(mm)
        if audio is None:
            print(f"  [req {i}] No waveform in multimodal_output.")
            continue
        out_path = str(output_dir / f"output_{i}.wav")
        save_audio(audio.cpu(), out_path, sr)

    print("Done.")


def parse_args():
    parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
    parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
    parser.add_argument("--text", default="你好，这是IndexTTS2语音合成测试。")
    parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
    parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
    parser.add_argument("--emo-text", default=None, help="Emotion description text.")
    parser.add_argument(
        "--emo-vector",
        type=float,
        nargs=8,
        default=None,
        help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
    )
    parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
    parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
    parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument(
        "--output-dir",
        default=os.path.join(
            os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
            "indextts2_output",
        ),
    )
    parser.add_argument("--deploy-config", default=None)
    parser.add_argument("--stage-init-timeout", type=int, default=600)
    return parser.parse_args()


if __name__ == "__main__":
    main(parse_args())