gepard_inference/modeling.py

"""Gepard multihead model — inference-only build.

Architecture:
  - Stock Qwen3.5 transformer backbone (no modifications)
  - N audio embedding tables (one per FSQ channel) → concat → GELU MLP →
    single frame embedding
  - N codebook output heads with per-channel vocabulary sizes
  - 1 binary stop head (end-of-speech prediction)
  - Optional voice-cloning stack: ``RefCompressor`` (codec ``ref_codes`` →
    K speaker prefix tokens) + learnable ``null_prefix``
  - No lm_head, no text generation, no training losses

This is a lean extraction of the training project's ``GepardModel``: the
module tree (and therefore the ``state_dict`` layout) matches published
checkpoints exactly, but every training-only code path (losses, diagnostics,
CFG dropout, gradient checkpointing plumbing) is removed. Generation is
driven externally by :class:`gepard_inference.runner.GepardRunner`.
"""

from __future__ import annotations

from types import SimpleNamespace
from typing import Dict, List, Optional

import torch
import torch.nn as nn

from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig

from .configuration import GepardConfig, reconcile_backbone_config
from .ref_compressor import RefCompressor

# Aliases people actually write in configs → canonical torch dtypes.
_DTYPE_MAP = {
    "bfloat16": torch.bfloat16,
    "bf16": torch.bfloat16,
    "float16": torch.float16,
    "fp16": torch.float16,
    "half": torch.float16,
    "float32": torch.float32,
    "fp32": torch.float32,
    "float": torch.float32,
}


def resolve_dtype(name) -> torch.dtype:
    """Map a config dtype string (or a torch.dtype) to the torch dtype."""
    if isinstance(name, torch.dtype):
        return name
    try:
        return _DTYPE_MAP[str(name).lower().removeprefix("torch.")]
    except KeyError:
        raise ValueError(
            f"unknown model dtype {name!r}; expected one of {sorted(_DTYPE_MAP)}"
        ) from None


class SupConProjectionHead(nn.Module):
    """2-layer MLP projection head from compressor d_model to a SupCon embedding.

    Training-only module; declared here purely for state-dict compatibility
    with checkpoints trained with the supervised-contrastive objective (its
    weights are loaded but never used at inference).
    """

    def __init__(self, d_model: int, hidden_dim: int = 128, projection_dim: int = 128):
        super().__init__()
        self.fc1 = nn.Linear(d_model, hidden_dim, bias=True)
        self.fc2 = nn.Linear(hidden_dim, projection_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(torch.relu(self.fc1(x)))


class GepardModel(nn.Module):
    """Multihead TTS model: Qwen3.5 backbone + N codebook heads + 1 stop head.

    Inference surface (used by the runner):
      - ``self.model`` — stock ``Qwen3_5TextModel`` backbone (KV-cached forward)
      - ``self._embed_audio`` — audio frame → backbone embedding
      - ``self.codebook_heads`` / ``self.vocab_sizes`` — per-channel logits
      - ``self.stop_head`` — end-of-speech logit
      - ``self.ref_compressor`` — voice-cloning prefix (may be None)
    """

    def __init__(
        self,
        config: Qwen3_5TextConfig,
        audio_heads: Dict[str, int],
        vc_config: Optional[SimpleNamespace] = None,
        codec: Optional[SimpleNamespace] = None,
        audio_embed_dim: int = 32,
    ):
        super().__init__()
        self.config = config
        self.audio_heads = audio_heads
        self.channel_names: List[str] = list(audio_heads.keys())
        self.vocab_sizes: List[int] = list(audio_heads.values())
        self.num_codebook_heads: int = len(audio_heads)
        hidden_size = config.hidden_size

        # Stock Qwen3.5 backbone
        self.model = Qwen3_5TextModel(config)

        # Audio input: per-codebook embeddings → concat → 2-layer GELU MLP.
        self.audio_embed_dim = audio_embed_dim
        self.audio_embeddings = nn.ModuleList([
            nn.Embedding(vocab_size, self.audio_embed_dim)
            for vocab_size in self.vocab_sizes
        ])
        self.audio_embed_proj = nn.Sequential(
            nn.Linear(self.num_codebook_heads * self.audio_embed_dim, hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, hidden_size),
            # Affine-free LayerNorm pins the frame-embedding scale.
            nn.LayerNorm(hidden_size, elementwise_affine=False),
        )
        # Rescales the unit-norm projection output to the text-embedding std.
        # Frozen parameter of shape [1]; the value comes from the checkpoint.
        self.audio_embed_scale = nn.Parameter(torch.tensor([1.0]), requires_grad=False)

        # Output heads: one Linear per channel, each with its own vocab size
        self.codebook_heads = nn.ModuleList([
            nn.Linear(hidden_size, vocab_size)
            for vocab_size in self.vocab_sizes
        ])
        self.stop_head = nn.Linear(hidden_size, 1)

        # Voice cloning: reference compressor + learnable null prefix (the
        # null prefix is a training-time CFG artifact kept for state-dict
        # compatibility; inference CFG contrasts text, not the speaker prefix).
        self.ref_compressor: Optional[RefCompressor] = None
        self.null_prefix: Optional[nn.Parameter] = None
        self.supcon_head: Optional[SupConProjectionHead] = None
        if vc_config is not None and vc_config.enabled:
            if codec is None:
                raise ValueError(
                    "GepardModel: vc_config.enabled=True but codec is None; "
                    "voice cloning needs the codec geometry."
                )
            self.ref_compressor = RefCompressor(
                codec=codec,
                compressor_cfg=vc_config.compressor,
                backbone_hidden_size=hidden_size,
            )
            self.null_prefix = nn.Parameter(
                torch.zeros(int(vc_config.compressor.num_queries), self.ref_compressor.d_model)
            )
            supcon = getattr(vc_config, "supcon", None)
            if supcon is not None and supcon.use_projection:
                self.supcon_head = SupConProjectionHead(
                    d_model=self.ref_compressor.d_model,
                    hidden_dim=int(supcon.projection_hidden_dim),
                    projection_dim=int(supcon.projection_dim),
                )

    def _embed_audio(self, audio_tensors: List[torch.LongTensor]) -> torch.FloatTensor:
        """Embed N codebook channels → concat → MLP → single frame embedding."""
        per_channel = [
            self.audio_embeddings[i](audio_tensors[i])
            for i in range(self.num_codebook_heads)
        ]
        concat = torch.cat(per_channel, dim=-1)            # [B, T_audio, N * audio_embed_dim]
        frame = self.audio_embed_proj(concat)              # [B, T_audio, d], unit-norm (LayerNorm)
        return frame * self.audio_embed_scale.to(frame.dtype)


def build_model(
    cfg: GepardConfig,
    attn_implementation: Optional[str] = None,
) -> GepardModel:
    """Reconstruct the model from a ``GepardConfig`` (fresh weights; matching shapes).

    Rebuilds the backbone from the nested config (reconciling
    ``partial_rotary_factor``). ``attn_implementation`` overrides the
    serialized backbone setting (inference typically wants ``eager`` where
    training used ``flash_attention_2``).
    """
    backbone_kwargs = dict(cfg.backbone_config)
    if attn_implementation is not None:
        # Drop any serialized attn setting so the override wins regardless of
        # which key form the transformers version wrote into to_dict().
        for k in ("attn_implementation", "_attn_implementation", "_attn_implementation_internal"):
            backbone_kwargs.pop(k, None)
        backbone_kwargs["attn_implementation"] = attn_implementation
    backbone = Qwen3_5TextConfig(**backbone_kwargs)
    reconcile_backbone_config(backbone, cfg)

    vc_config = None
    codec = None
    if cfg.vc_enabled:
        comp = cfg.voice_cloning.get("compressor", {})
        sc = (cfg.voice_cloning.get("training") or {}).get("supcon", {})
        vc_config = SimpleNamespace(
            enabled=True,
            compressor=SimpleNamespace(
                num_queries=int(comp["num_queries"]),
                num_layers=int(comp["num_layers"]),
                num_heads=int(comp["num_heads"]),
                d_model=comp.get("d_model"),
                ffn_hidden_size_multiplier=int(comp["ffn_hidden_size_multiplier"]),
                dropout=float(comp.get("dropout", 0.1)),
                queries_init_std=0.02,
            ),
            supcon=SimpleNamespace(
                use_projection=bool(sc.get("enabled", False)) and bool(sc.get("use_projection", False)),
                projection_hidden_dim=int(sc.get("projection_hidden_dim", 128)),
                projection_dim=int(sc.get("projection_dim", 128)),
            ),
        )
        # Duck-typed stand-in for the codec group: the compressor only reads
        # num_layers / fsq_levels / do_unfold.
        codec = SimpleNamespace(
            num_layers=int(cfg.codec["num_layers"]),
            fsq_levels=list(cfg.codec["fsq_levels"]),
            do_unfold=bool(cfg.codec.get("do_unfold", True)),
        )

    return GepardModel(
        config=backbone,
        audio_heads=dict(cfg.audio_heads),
        vc_config=vc_config,
        codec=codec,
        audio_embed_dim=cfg.audio_embed_dim,
    )