#!/usr/bin/env python3
"""Audio8 TTS Preview 0.6B backend for ZeroGPU.

DualAR (Fish-Audio-S2-Pro-style) 0.6B multilingual TTS with zero-shot voice
cloning: https://huggingface.co/Audio8/Audio8-TTS-Preview-0.6b. Requires
trust_remote_code=True (custom `arktts` modeling/processing code shipped in
the model repo).

Requires transformers<5 (see requirements.txt) — the model's hand-rolled
generate() loop produces near-silent, never-terminating output under
transformers 5.x. Confirmed by isolating the variable on a local CUDA GPU:
same code, same hardware, only the transformers version changed between a
working run (4.57.6) and a broken one (5.15.0). The Gradio SDK is pinned to
5.50.0 in README.md specifically so transformers<5 and gradio's own
huggingface_hub requirement can both be satisfied in one pip resolution.

``load()`` runs once at app startup and moves the model to "cuda" eagerly —
ZeroGPU's CUDA emulation packs the tensors to disk and streams them into VRAM
on the first real request, so lazy loading inside the decorated handler would
cost every user instead of only the first.
"""
import logging
import os

import torch

MODEL_REPO = "Audio8/Audio8-TTS-Preview-0.6b"

_processor = None
_model = None


def load():
    """Load the processor + model onto cuda. Idempotent."""
    global _processor, _model
    if _model is not None:
        return

    from transformers import AutoModel, AutoProcessor

    logging.info(f"Loading Audio8 TTS ({MODEL_REPO})…")
    token = os.environ.get("HF_TOKEN")
    _processor = AutoProcessor.from_pretrained(
        MODEL_REPO, token=token, trust_remote_code=True
    )
    _model = (
        AutoModel.from_pretrained(
            MODEL_REPO, token=token, trust_remote_code=True, dtype=torch.bfloat16
        )
        .eval()
        .to("cuda")
    )
    logging.info("Audio8 TTS ready.")


def generate(text, voice_ref=None, reference_text=None, temperature=0.8,
             top_p=0.95, top_k=50, max_new_tokens=1024, seed=-1):
    """Generate speech with Audio8 TTS Preview.

    voice_ref: optional path to a reference clip for zero-shot cloning.
    reference_text: transcript of voice_ref. The processor *requires* a
        non-empty reference_text whenever voice_ref is given — the model
        conditions generation on the text/audio alignment, not just the
        audio. Raises ValueError if missing.

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

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

    call_kwargs = {}
    if voice_ref:
        if not reference_text or not reference_text.strip():
            raise ValueError(
                "This model needs a transcript of the reference clip to clone it. "
                "Fill in \"Reference transcript\" (auto-transcription may have failed) "
                "or clear the reference audio to generate without cloning."
            )
        call_kwargs["reference_audio"] = [voice_ref]
        call_kwargs["reference_text"] = [reference_text.strip()]

    inputs = _processor(text=[text], return_tensors="pt", **call_kwargs)
    inputs = {k: v.to("cuda") for k, v in inputs.items()}

    with torch.inference_mode():
        output = _model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            do_sample=True,
            return_dict_in_generate=True,
        )
        waveforms, waveform_lengths = _model.decode_audio(output.codes)

    audio = waveforms[0, : int(waveform_lengths[0])].float().cpu().numpy()
    if audio.size == 0:
        raise RuntimeError("Audio8 TTS produced no audio — try again or adjust the text.")
    return audio, _model.config.codec_sample_rate