#!/usr/bin/env python3
"""Higgs Audio v3 (4B) TTS backend for ZeroGPU.

Ported from the `transformers`-native demo at
multimodalart/higgs-audio-v3-tts, which wraps bosonai/higgs-audio-v3-tts-4b
via multimodalart/higgs-audio-v3-tts-4b-transformers.

``load()`` is called once at app startup. ZeroGPU's CUDA emulation packs the
model tensors before a real GPU is assigned to the decorated request handler.
"""
import logging
import os

import torch

MODEL_REPO = "multimodalart/higgs-audio-v3-tts-4b-transformers"

_tokenizer = None
_model = None
_sample_rate = None


def load():
    """Load the tokenizer, model, and audio codec onto cuda. Idempotent."""
    global _tokenizer, _model, _sample_rate
    if _model is not None:
        return

    from transformers import AutoModelForCausalLM, AutoTokenizer

    logging.info(f"Loading Higgs Audio v3 ({MODEL_REPO})…")
    token = os.environ.get("HF_TOKEN")
    _tokenizer = AutoTokenizer.from_pretrained(
        MODEL_REPO,
        token=token,
        trust_remote_code=True,
    )
    _model = (
        AutoModelForCausalLM.from_pretrained(
            MODEL_REPO,
            token=token,
            trust_remote_code=True,
            dtype=torch.bfloat16,
        )
        .to("cuda")
        .eval()
    )
    _model.get_audio_codec()  # preload the 24 kHz codec
    _sample_rate = _model.config.sample_rate
    logging.info("Higgs Audio v3 ready.")


def generate(text, voice_ref=None, reference_text=None, temperature=0.7,
             top_p=0.95, top_k=50, max_new_tokens=2048, seed=-1):
    """Generate speech with Higgs Audio v3.

    voice_ref: optional path to a reference clip for zero-shot cloning.
    reference_text: optional transcript of voice_ref — improves cloning when
        provided, but generation works without it.

    Returns (waveform: 1-D numpy array, sample_rate: int).
    """
    if _model is None:
        raise RuntimeError("Higgs Audio v3 is not loaded — call higgs_backend.load() at startup.")

    import soundfile as sf

    if seed is not None and int(seed) >= 0:
        torch.manual_seed(int(seed))

    kwargs = dict(
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        top_p=float(top_p) if float(top_p) < 1.0 else None,
        top_k=int(top_k) if int(top_k) > 0 else None,
    )

    if voice_ref:
        data, sr = sf.read(voice_ref, dtype="float32", always_2d=True)  # [L, C]
        wav = torch.from_numpy(data).mean(dim=1)  # to mono [L]
        kwargs["reference_audio"] = wav
        kwargs["reference_sample_rate"] = sr
        if reference_text and reference_text.strip():
            kwargs["reference_text"] = reference_text.strip()

    audio = _model.generate_speech(text, _tokenizer, **kwargs)
    if audio.numel() == 0:
        raise RuntimeError("Higgs Audio v3 produced no audio — try again or adjust the text.")
    return audio.numpy(), _sample_rate