#!/usr/bin/env python3
"""Synthesize speech with TalkingFlower TTS.

Usage:
    python inference.py "你好，我是会说话的花朵。" output.wav

Requirements:
    pip install onnxruntime torch torchaudio transformers \
                hyperpyyaml soundfile librosa einops

This bundle is self-contained: the cosyvoice package and Matcha-TTS library
are included under cosyvoice/ and third_party/Matcha-TTS/ respectively.
No separate CosyVoice repository clone is needed.

This script targets SFT-only inference (text → speech for the bundled speaker).
For zero-shot, cross-lingual, prompt-audio, or TRAINING, see the
"Restoring full functionality" section below.
"""
import os
import sys
os.environ.setdefault("COSYVOICE_DISABLE_TEXT_FRONTEND", "1")
os.environ.setdefault("COSYVOICE_LLM_EMBED_INT8", "1")

import onnxruntime
import torch
import torchaudio
import transformers

# ---------------------------------------------------------------------------
# Path setup — cosyvoice package and Matcha-TTS are bundled alongside this file
# ---------------------------------------------------------------------------
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
sys.path.insert(0, os.path.join(_HERE, "third_party/Matcha-TTS"))


# ===========================================================================
# Slim-bundle monkey-patch  ↓↓↓  (remove this block to restore full upstream)
# ===========================================================================
#
# This HF bundle ships WITHOUT the two ONNX models that CosyVoice's frontend
# normally loads at init:
#
#     campplus.onnx              (27 MB)   speaker encoder
#     speech_tokenizer_v3.onnx   (925 MB)  audio → discrete speech tokens
#
# For SFT mode (this script's purpose) both sessions are created but NEVER
# called — the speaker embedding comes from spk2info.pt and there is no prompt
# audio to tokenize. We save ~952 MB by skipping the InferenceSession() calls
# when the file is absent.
#
# ─── Restoring full functionality ──────────────────────────────────────────
# If you want zero-shot, cross-lingual, prompt-audio inference, or to RESUME
# TRAINING, you need both ONNX files. To restore:
#   1. Download from upstream:
#        huggingface-cli download FunAudioLLM/CosyVoice2-0.5B \
#            campplus.onnx speech_tokenizer_v2.onnx \
#            --local-dir <this-bundle-dir>
#      (CosyVoice2 ships speech_tokenizer_v2; CosyVoice3 uses v3 — get v3 from
#       the Fun-CosyVoice3-0.5B repo on ModelScope, file
#       `speech_tokenizer_v3.onnx`.)
#   2. Delete or comment out the patch block below.
# ───────────────────────────────────────────────────────────────────────────

_original_inference_session = onnxruntime.InferenceSession


def _maybe_inference_session(path_or_bytes, *args, **kwargs):
    """Return a real ORT session, or None if the model file is missing."""
    if isinstance(path_or_bytes, str) and not os.path.exists(path_or_bytes):
        return None
    return _original_inference_session(path_or_bytes, *args, **kwargs)


# Patch BEFORE importing CosyVoice — the frontend's __init__ calls
# onnxruntime.InferenceSession() during model construction.
onnxruntime.InferenceSession = _maybe_inference_session


# ─── Patch 2: skip Qwen2 base-LLM weight loading (saves 988 MB) ────────────
#
# The bundled `CosyVoice-BlankEN/` directory ships ONLY the small files needed
# to construct the Qwen2-0.5B architecture and tokenizer:
#
#     config.json, generation_config.json, tokenizer_config.json,
#     vocab.json, merges.txt
#
# The 988 MB `model.safetensors` (base Qwen2 pretrained weights) is OMITTED
# because the load path is:
#     1. Qwen2ForCausalLM.from_pretrained(<this-dir>)   ← would load 988 MB
#     2. llm.load_state_dict(torch.load('llm.pt'))      ← immediately overwrites
# Step 1's weights are thrown away on step 2, so we can skip them.
#
# This patch detects an absent safetensors and instantiates the Qwen2 model
# from config only — letting `llm.pt` populate every parameter.
#
# ─── Restoring full functionality (training, finetuning from this bundle) ──
# If you want to RESUME TRAINING from this checkpoint:
#   1. Download the base Qwen2 weights into CosyVoice-BlankEN/:
#        huggingface-cli download FunAudioLLM/CosyVoice2-0.5B \
#            CosyVoice-BlankEN/model.safetensors \
#            --local-dir <this-bundle-dir>
#   2. Delete or comment out this Patch 2 block — the upstream from_pretrained
#      will then load the base weights as usual.
# (Inference is unaffected either way; the patch is also harmless when the
# safetensors file IS present.)
# ───────────────────────────────────────────────────────────────────────────

_original_qwen2_from_pretrained = transformers.Qwen2ForCausalLM.from_pretrained


def _qwen2_from_pretrained_or_config(pretrained_model_name_or_path, *args, **kwargs):
    """Use Qwen2Config + Qwen2ForCausalLM(config) when the weights file is absent."""
    weight_files = ('model.safetensors', 'pytorch_model.bin',
                    'model.safetensors.index.json', 'pytorch_model.bin.index.json')
    if (isinstance(pretrained_model_name_or_path, str)
            and os.path.isdir(pretrained_model_name_or_path)
            and not any(os.path.exists(os.path.join(pretrained_model_name_or_path, f))
                        for f in weight_files)):
        config = transformers.Qwen2Config.from_pretrained(pretrained_model_name_or_path)
        return transformers.Qwen2ForCausalLM(config)
    return _original_qwen2_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)


transformers.Qwen2ForCausalLM.from_pretrained = _qwen2_from_pretrained_or_config

# ===========================================================================
# End of slim-bundle monkey-patches  ↑↑↑
# ===========================================================================


from cosyvoice.cli.cosyvoice import CosyVoice3  # noqa: E402

MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
INSTRUCT = "You are a helpful assistant.<|endofprompt|>"
SPK_ID = "TalkingFlower"


def remove_tail_click(audio, sr, search_s=0.20, burst_thresh=0.05,
                      silence_thresh=0.02, win_ms=5, fade_ms=3):
    """Remove LLM-induced high-energy burst at the tail of the waveform.

    The fine-tuned LLM occasionally emits a final-token sequence that the
    vocoder renders as a short impulsive transient. Pattern is always:
    [speech] → [silence 30-40ms] → [burst 10-25ms] → [EOF]. We find the gap
    of silence before the burst and zero out from there, with a 3ms fade in.
    Clean clips (last-window RMS ≈ 0) pass through unchanged.
    """
    ch = audio[0]
    win_n = int(sr * win_ms / 1000)
    search_n = min(int(sr * search_s), ch.shape[0])
    fade_n = int(sr * fade_ms / 1000)
    tail = ch[-search_n:]
    n_wins = search_n // win_n
    rms = [tail[i * win_n:(i + 1) * win_n].pow(2).mean().sqrt().item()
           for i in range(n_wins)]
    if rms[-1] < burst_thresh:
        return audio
    cut_win = next((i for i in range(n_wins - 2, -1, -1)
                    if rms[i] < silence_thresh), None)
    if cut_win is None:
        return audio
    cut = ch.shape[0] - search_n + cut_win * win_n
    out = audio.clone()
    out[0, cut:] = 0.0
    if fade_n > 0 and cut >= fade_n:
        out[0, cut - fade_n:cut] *= torch.linspace(1.0, 0.0, fade_n)
    return out


def synthesize(text, out_path):
    model = CosyVoice3(MODEL_DIR, fp16=True)
    for output in model.inference_sft(
        INSTRUCT + text,
        spk_id=SPK_ID,
        stream=False,
        text_frontend=False,
    ):
        audio = remove_tail_click(output["tts_speech"], model.sample_rate)
        import soundfile as sf
        sf.write(out_path, audio.squeeze(0).cpu().numpy(), model.sample_rate)
        print(f"Saved {out_path}  ({audio.shape[1] / model.sample_rate:.2f}s)")
        break


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    text = sys.argv[1]
    out = sys.argv[2] if len(sys.argv) > 2 else "output.wav"
    synthesize(text, out)