gepard_inference/runner.py

"""Gepard autoregressive TTS runner with text classifier-free guidance.

``GepardRunner`` is the canonical inference entry point: with the default
``cfg_scale=1.0`` it runs plain single-pass generation; raising ``cfg_scale``
turns on text classifier-free guidance.

Usage:
    runner = GepardRunner.from_checkpoint("nineninesix/kani3-dpo-cfg-round-2")
    tokens = runner.generate("Hello world", ref_codes=ref_codes,
                             temperature=0.3, cfg_scale=3.0)
    # tokens: (num_heads, T) — ready for UnfoldedCodecModel.decode_from_codes()

Why text-CFG
------------
On short inputs the K speaker-prefix tokens dominate the hidden state and the
1–2 text tokens drown → the model never "latches" and runs away. CFG runs two
forward passes per frame:

    cond   = [ prefix | SOT text EOT SOS | audio... ]   (full conditioning)
    uncond = [ prefix | SOT      EOT SOS | audio... ]   (SAME prefix, NO text)

and guides the per-head logits toward what the text adds:

    logit_guided = logit_uncond + cfg_scale * (logit_cond - logit_uncond)

The speaker prior (common to both branches) cancels in the difference, so the
text-specific direction is amplified. Cost: 2x forward per guided frame;
onset-only guidance (``cfg_frames``) keeps it to the first N frames.
"""

from typing import List, Optional

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DynamicCache
from safetensors.torch import load_file as load_safetensors

from .checkpoint_io import normalize_scalar_shapes, resolve_safetensors
from .configuration import load_gepard_config
from .modeling import GepardModel, build_model, resolve_dtype
from .text_repetition import TextRepetitionConfig, TextRepeater


class FullAttnCache(Qwen3_5DynamicCache):
    """Qwen3_5DynamicCache that works for full-attention-only models
    (no linear_attention layers)."""

    def __init__(self, config):
        # Skip parent __init__ — it crashes when layer_types has no "linear_attention"
        self.layer_types = config.layer_types
        self.transformer_layers = [
            i for i in range(config.num_hidden_layers)
            if self.layer_types[i] == "full_attention"
        ]
        self.last_linear_layer = -1  # no linear layers in this model
        self.conv_states = [None for _ in range(config.num_hidden_layers)]
        self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
        self.key_cache = [None for _ in range(config.num_hidden_layers)]
        self.value_cache = [None for _ in range(config.num_hidden_layers)]


class GepardRunner:
    """Autoregressive inference for GepardModel with KV cache and text-CFG.

    Generates per-channel FSQ token sequences from text. The returned tensor
    shape (num_heads, T) is directly compatible with
    ``UnfoldedCodecModel.decode_from_codes()`` after unsqueezing the batch dim.

    Token format follows the training layout:
        text_ids = [BOS_text, ...text tokens..., EOT, BOS_audio]
        then audio frames are generated autoregressively.
    """

    def __init__(
        self,
        model: GepardModel,
        tokenizer,
        device: torch.device,
        repetition: Optional[TextRepetitionConfig] = None,
        special_tokens: Optional[dict] = None,
    ):
        self.model = model.eval()
        self.tokenizer = tokenizer
        self.device = torch.device(device)
        self.channel_names: List[str] = model.channel_names
        self.num_heads: int = model.num_codebook_heads
        # Key names follow the checkpoint's GepardConfig.special_tokens convention.
        st = special_tokens or {}
        self.BOS_TEXT = int(st["start_of_text"])
        self.EOT = int(st["end_of_text"])
        self.BOS_AUDIO = int(st["start_of_speech"])
        # Adaptive text repetition — must mirror the training-time layout
        # exactly (same special ids, same target/threshold).
        self.repetition_cfg = repetition or TextRepetitionConfig()
        self.repeater = TextRepeater(
            self.repetition_cfg, self.BOS_TEXT, self.EOT, self.BOS_AUDIO,
        )

    # ------------------------------------------------------------------
    # Factory
    # ------------------------------------------------------------------

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: str,
        device: Optional[str] = None,
        attn_implementation: str = "eager",
    ) -> "GepardRunner":
        """Load model from a self-describing checkpoint (local dir or HF repo id).

        The checkpoint must carry a ``gepard_config.json``; the runner rebuilds
        the exact architecture (backbone, audio heads, voice-cloning
        compressor) from the checkpoint alone — no training configs needed.

        Args:
            checkpoint_path: HF-style checkpoint (contains model.safetensors +
                gepard_config.json + tokenizer files).
            device: 'cuda', 'cpu', or None (auto-detect).
            attn_implementation: 'eager' or 'flash_attention_2'.
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device)

        gepard_cfg = load_gepard_config(checkpoint_path)
        if gepard_cfg is None:
            raise FileNotFoundError(
                f"{checkpoint_path!r} has no gepard_config.json — only "
                "self-describing checkpoints are supported by this package."
            )
        model = build_model(gepard_cfg, attn_implementation=attn_implementation)
        model_dtype = resolve_dtype(gepard_cfg.model_dtype)
        print(f"[{cls.__name__}] model config: {checkpoint_path}/gepard_config.json")

        # Load full state dict — covers backbone AND custom heads in one shot.
        sf_path = resolve_safetensors(checkpoint_path)
        state_dict = load_safetensors(sf_path, device="cpu")
        reshaped = normalize_scalar_shapes(state_dict, model)
        if reshaped:
            print(f"[{cls.__name__}] legacy 0-dim params reshaped to match model: {reshaped}")
        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        if missing:
            print(f"[{cls.__name__}] missing keys ({len(missing)}): {missing[:3]} ...")
        if unexpected:
            print(f"[{cls.__name__}] unexpected keys ({len(unexpected)}): {unexpected[:3]} ...")

        model = model.to(device=device, dtype=model_dtype)

        tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
        repetition = TextRepetitionConfig.from_config(dict(gepard_cfg.text_repetition))

        return cls(
            model=model,
            tokenizer=tokenizer,
            device=device,
            repetition=repetition,
            special_tokens=dict(gepard_cfg.special_tokens),
        )

    # ------------------------------------------------------------------
    # Core generate
    # ------------------------------------------------------------------

    @torch.no_grad()
    def generate(
        self,
        text: str,
        ref_codes: Optional[torch.Tensor] = None,   # [1, T_ref, C_total] long — unfolded codec codes
        ref_mask: Optional[torch.Tensor] = None,    # [1, T_ref] bool, True = real frame
        temperature: float = 1.0,
        top_k: int = 0,
        stop_threshold: float = 0.5,
        max_frames: int = 2000,
        repetition_penalty: float = 1.0,
        repetition_window: int = 32,
        force_stop_frames: Optional[int] = None,
        cfg_scale: float = 1.0,
        cfg_frames: Optional[int] = None,
        cfg_max_text_tokens: Optional[int] = None,
        cfg_uncond_mode: str = "empty_text",
    ) -> torch.LongTensor:
        """Generate audio tokens for ``text``, optionally with a reference voice
        and text classifier-free guidance.

        Args:
            text: Input text to synthesize.
            ref_codes: Unfolded codec codes for the reference voice, shape
                [1, T_ref, C_total]. Ignored if the model has no ref_compressor.
            ref_mask: Boolean mask for ref_codes, True = real frame. If None
                and ref_codes is given, all frames are assumed real.
            temperature: Sampling temperature (1.0 = no change).
            top_k: If > 0, keep only top-k logits before sampling.
            stop_threshold: Sigmoid threshold on the stop head to terminate.
            max_frames: Hard cap on the number of generated audio frames.
            repetition_penalty: >1.0 penalises tokens seen in the recent window
                per head. 1.0 = disabled. Typical: 1.1–1.3.
            repetition_window: Number of recent frames tracked for
                repetition_penalty. 0 = all history.
            force_stop_frames: Deterministic guardrail — hard-stop at this many
                frames regardless of the stop head. None = disabled.
            cfg_scale: Guidance weight w. 1.0 = disabled (plain single-pass).
                2.0–3.0 typical; higher = stronger text emphasis but risks
                artefacts / premature stop.
            cfg_frames: Onset-only guidance. If set, CFG is applied only to the
                first N frames; later frames use the cond branch alone (the
                uncond branch is then skipped to save compute). None = guide
                every frame.
            cfg_max_text_tokens: Length gate for text-CFG. If set and the input
                has MORE than this many text tokens, CFG is turned off
                (``cfg_scale`` forced to 1.0) for this call. CFG rescues short
                phrases from running away, but on longer text it rushes
                delivery (~2x) and flattens prosody — so it is only worth
                applying to short inputs. None = no gate (always honour
                ``cfg_scale``).
            cfg_uncond_mode: How to build the text-free uncond branch:
                "empty_text" → [BOS_text, EOT, BOS_audio] (cleanest contrast);
                "audio_only" → [BOS_audio] (most aggressive; more OOD).

        Returns:
            LongTensor of shape (num_heads, T) — per-channel FSQ codes,
            compatible with ``UnfoldedCodecModel.decode_from_codes()`` after
            adding a batch dim: ``tokens.unsqueeze(0)``.
        """
        # 1. Build cond text_ids (with adaptive repetition).
        text_token_ids = self.tokenizer.encode(text, add_special_tokens=False)

        # Length gate: CFG only helps short phrases latch onto the speech
        # manifold. On longer text it rushes delivery (~2x) and flattens
        # prosody, so disable it past the threshold.
        if cfg_max_text_tokens is not None and len(text_token_ids) > cfg_max_text_tokens:
            if cfg_scale != 1.0:
                print(
                    f"[GepardRunner] text has {len(text_token_ids)} tokens "
                    f"(> cfg_max_text_tokens={cfg_max_text_tokens}); disabling CFG."
                )
            cfg_scale = 1.0

        cfg_on = cfg_scale != 1.0
        cond_ids = self.repeater.expand(text_token_ids)

        # 1b. Build uncond text_ids (SAME prefix later, text removed).
        if cfg_uncond_mode == "empty_text":
            uncond_ids = [self.BOS_TEXT, self.EOT, self.BOS_AUDIO]
        elif cfg_uncond_mode == "audio_only":
            uncond_ids = [self.BOS_AUDIO]
        else:
            raise ValueError(f"unknown cfg_uncond_mode={cfg_uncond_mode!r}")

        # 2. Voice-cloning prefix (shared by both branches).
        prefix_embeds = self._compute_ref_prefix(ref_codes, ref_mask)

        # 3. Prefill cond (and uncond if CFG is on).
        cond_hidden, cond_cache, K, T_text_cond = self._prefill(prefix_embeds, cond_ids)
        if cfg_on:
            unc_hidden, unc_cache, _, T_text_unc = self._prefill(prefix_embeds, uncond_ids)
        else:
            unc_hidden = unc_cache = None
            T_text_unc = 0

        # 4. First audio frame from the BOS_audio position of each branch.
        guide_now = cfg_on and (cfg_frames is None or 0 < cfg_frames)
        first_frame = self._sample_frame(
            cond_hidden[:, -1:, :],
            unc_hidden[:, -1:, :] if guide_now else None,
            cfg_scale if guide_now else 1.0,
            temperature, top_k, repetition_penalty, None,
        )
        generated: List[torch.LongTensor] = [first_frame]

        # 5. Autoregressive loop.
        for step in range(1, max_frames):
            prev_frame = generated[-1]
            frame_embed = self._embed_frame(prev_frame.unsqueeze(0))  # (1, 1, d)

            # Attention mask length == total KV after this frame is appended:
            # prefix + text + (step-1) cached audio frames + current frame.
            cond_hidden, cond_cache = self._decode_step(
                frame_embed, cond_cache, K + T_text_cond + step,
            )
            guide_now = cfg_on and (cfg_frames is None or step < cfg_frames)
            if cfg_on and guide_now:
                unc_hidden, unc_cache = self._decode_step(
                    frame_embed, unc_cache, K + T_text_unc + step,
                )
            # Past the onset window guidance is off for good, so the uncond
            # forward is skipped entirely to save compute.

            # Deterministic guardrail: hard-stop runaway the stop head misses.
            if force_stop_frames is not None and len(generated) >= force_stop_frames:
                break

            # Stop decision from the COND branch (the real conditioned model).
            stop_logit = self.model.stop_head(cond_hidden[:, -1, :])  # (1, 1)
            stop_prob = torch.sigmoid(stop_logit.squeeze()).item()
            if stop_prob > stop_threshold:
                break

            if repetition_penalty != 1.0:
                window = generated if repetition_window == 0 else generated[-repetition_window:]
            else:
                window = None

            next_frame = self._sample_frame(
                cond_hidden,
                unc_hidden if (cfg_on and guide_now) else None,
                cfg_scale if (cfg_on and guide_now) else 1.0,
                temperature, top_k, repetition_penalty, window,
            )
            generated.append(next_frame)

        # 6. Stack frames: list of (num_heads,) → (num_heads, T)
        tokens = torch.stack(generated, dim=0).T.contiguous()
        return tokens

    # ------------------------------------------------------------------
    # Forward helpers
    # ------------------------------------------------------------------

    def _prefill(self, prefix_embeds, input_ids):
        """Run the prefill pass for one branch.

        Returns (hidden, past_key_values, K, T_text) where hidden is the full
        prefilled hidden state and the cache covers K + T_text positions.
        """
        text_ids = torch.tensor([input_ids], dtype=torch.long, device=self.device)  # (1, T_text)
        T_text = text_ids.shape[1]
        text_embeds = self.model.model.embed_tokens(text_ids)  # (1, T_text, d)

        if prefix_embeds is not None:
            inputs_embeds = torch.cat([prefix_embeds, text_embeds], dim=1)  # (1, K+T_text, d)
            K = prefix_embeds.size(1)
        else:
            inputs_embeds = text_embeds
            K = 0

        attn_mask = torch.ones(1, K + T_text, dtype=torch.long, device=self.device)
        past_key_values = FullAttnCache(self.model.config)
        out = self.model.model(
            inputs_embeds=inputs_embeds,
            attention_mask=attn_mask,
            use_cache=True,
            past_key_values=past_key_values,
        )
        return out.last_hidden_state, out.past_key_values, K, T_text

    def _decode_step(self, frame_embed, past_key_values, kv_len):
        """One AR decode step for one branch. Returns (hidden, past_key_values)."""
        attn_mask = torch.ones(1, kv_len, dtype=torch.long, device=self.device)
        out = self.model.model(
            inputs_embeds=frame_embed,
            attention_mask=attn_mask,
            use_cache=True,
            past_key_values=past_key_values,
        )
        return out.last_hidden_state, out.past_key_values

    def _compute_ref_prefix(self, ref_codes, ref_mask):
        """Voice-cloning prefix (K speaker tokens) if ref_compressor is present."""
        if self.model.ref_compressor is None or ref_codes is None:
            return None
        ref_codes = ref_codes.to(self.device)
        if ref_mask is None:
            ref_mask = torch.ones(
                ref_codes.shape[0], ref_codes.shape[1],
                dtype=torch.bool, device=ref_codes.device,
            )
        prefix_embeds, _ = self.model.ref_compressor(ref_codes, ref_mask)  # [1, K, d]
        return prefix_embeds

    # ------------------------------------------------------------------
    # Sampling helpers
    # ------------------------------------------------------------------

    def _embed_frame(self, frame_tokens: torch.LongTensor) -> torch.FloatTensor:
        """Embed a single audio frame via the model's audio-embedding stack.

        Args:
            frame_tokens: (B, num_heads) — one token per channel
        Returns:
            (B, 1, d)
        """
        channel_tokens = [frame_tokens[:, i] for i in range(self.num_heads)]
        emb = self.model._embed_audio(channel_tokens)  # (B, d)
        return emb.unsqueeze(1)  # (B, 1, d)

    def _sample_head_logits(
        self,
        logits: torch.FloatTensor,
        head_index: int,
        vocab_size: int,
        temperature: float,
        top_k: int,
        repetition_penalty: float,
        recent_frames: Optional[List[torch.LongTensor]],
    ) -> torch.LongTensor:
        """Apply repetition penalty / temperature / top-k to one head's fp32
        logits and sample."""
        if repetition_penalty != 1.0 and recent_frames:
            seen = {t[head_index].item() for t in recent_frames}
            for tok in seen:
                if logits[0, tok] > 0:
                    logits[0, tok] /= repetition_penalty
                else:
                    logits[0, tok] *= repetition_penalty

        if temperature != 1.0:
            logits = logits / temperature

        if top_k > 0:
            k = min(top_k, vocab_size)
            topk_vals, _ = torch.topk(logits, k)
            threshold = topk_vals[:, -1:]
            logits = logits.masked_fill(logits < threshold, float("-inf"))

        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1).squeeze()  # (1,1) → scalar

    def _sample_frame(
        self,
        cond_hidden: torch.FloatTensor,
        uncond_hidden: Optional[torch.FloatTensor],
        cfg_scale: float,
        temperature: float,
        top_k: int,
        repetition_penalty: float = 1.0,
        recent_frames: Optional[List[torch.LongTensor]] = None,
    ) -> torch.LongTensor:
        """Sample one audio frame from all codebook heads independently.

        If ``uncond_hidden`` is given and ``cfg_scale != 1.0``, per-head logits
        are guided BEFORE temperature/top-k/sampling:
            logit = logit_uncond + cfg_scale * (logit_cond - logit_uncond)
        Guidance is applied in fp32 logit space, on raw head outputs.
        """
        h_c = cond_hidden[:, -1, :]  # (1, d)
        h_u = uncond_hidden[:, -1, :] if uncond_hidden is not None else None
        do_cfg = h_u is not None and cfg_scale != 1.0

        frame_tokens = []
        for i, (head, vocab_size) in enumerate(zip(self.model.codebook_heads, self.model.vocab_sizes)):
            logits = head(h_c).float()  # (1, vocab_size) — fp32 for sampling stability
            if do_cfg:
                logits_u = head(h_u).float()
                logits = logits_u + cfg_scale * (logits - logits_u)

            token = self._sample_head_logits(
                logits, i, vocab_size, temperature, top_k, repetition_penalty, recent_frames,
            )
            frame_tokens.append(token)

        return torch.stack(frame_tokens)  # (num_heads,)