"""Unified checkpoint loading and audio-to-semantic-token inference.""" import json import math from pathlib import Path import numpy as np import torch import torch.nn.functional as F from safetensors import safe_open from torch import nn from .encoder import MERT2Encoder, SAMPLE_RATE as ENCODER_RATE CODEC_SIZE = 32768 HEAD_WINDOW = 512 class SwiGLU(nn.Module): def __init__(self, width: int, hidden: int): super().__init__() self.norm = nn.LayerNorm(width) self.gate = nn.Linear(width, hidden, bias=False) self.up = nn.Linear(width, hidden, bias=False) self.down = nn.Linear(hidden, width, bias=False) def forward(self, x): hidden = self.norm(x) return x + self.down(F.silu(self.gate(hidden)) * self.up(hidden)) class CVQHead(nn.Module): def __init__(self, width: int, dim: int, hidden: int, tau: float): super().__init__() self.adapter = SwiGLU(width, hidden) if hidden else None self.proj = nn.Linear(width, dim) self.codebook = nn.Parameter(torch.empty(CODEC_SIZE, dim)) self.tau = tau def forward(self, x): if self.adapter is not None: x = self.adapter(x) return F.normalize(self.proj(x), dim=-1) @ F.normalize(self.codebook, dim=-1).T / self.tau def to_token_clock(per_frame: np.ndarray, offset: int) -> np.ndarray: """Delay predictions so feature frame t supplies token t + offset.""" indices = np.clip(np.arange(len(per_frame)) - offset, 0, len(per_frame) - 1) return per_frame[indices] class WannabeTokenizer(nn.Module): """Load a unified WannabeTokenizer safetensors file without remote code.""" def __init__(self, checkpoint: str | Path, device: str | torch.device = "cuda"): super().__init__() self.checkpoint = Path(checkpoint) if not self.checkpoint.is_file(): raise FileNotFoundError(f"Tokenizer checkpoint not found: {self.checkpoint}") self.device = torch.device(device) with safe_open(self.checkpoint, framework="pt") as f: metadata = dict(f.metadata() or {}) self._validate_format(metadata) config = self._json_object(metadata, "encoder.config") source = self._json_object(metadata, "head.source_spec") self.layer = self._integer(source, "layer", minimum=0) causal = source.get("causal") if type(causal) is not bool: raise ValueError("head.source_spec.causal must be a boolean") self.chunk = float(source.get("chunk", 0)) if not math.isfinite(self.chunk) or self.chunk < 0 or 0 < self.chunk < 1: raise ValueError("head.source_spec.chunk must be 0 or at least 1 second") self.offset = self._integer(metadata, "head.offset", minimum=0) self.instnorm = self._boolean(metadata, "head.instnorm") tau = float(metadata.get("head.tau", 0.07)) if not math.isfinite(tau) or tau <= 0: raise ValueError("head.tau must be a finite positive number") encoder_keys = [key for key in f.keys() if key.startswith("encoder.")] head_keys = [key for key in f.keys() if key.startswith("head.")] if not encoder_keys or not head_keys: raise ValueError("Unified checkpoint must contain encoder.* and head.* tensors") encoder_state = {key[8:]: f.get_tensor(key) for key in encoder_keys} head_state = {key[5:]: f.get_tensor(key).float() for key in head_keys} self.encoder = MERT2Encoder(config, causal=causal, layers=self.layer + 1) kept_encoder_state = { key: value for key, value in encoder_state.items() if not (key.startswith("layers.") and int(key.split(".")[1]) > self.layer) } self.encoder.load_state_dict(kept_encoder_state, strict=True) try: dim = head_state["codebook"].shape[1] width = head_state["proj.weight"].shape[1] hidden = head_state["adapter.up.weight"].shape[0] if "adapter.up.weight" in head_state else 0 except (KeyError, IndexError) as exc: raise ValueError("Checkpoint is missing valid CVQ head tensors") from exc if head_state["codebook"].shape[0] != CODEC_SIZE: raise ValueError(f"Expected {CODEC_SIZE} codebook entries") self.head = CVQHead(width, dim, hidden, tau) self.head.load_state_dict(head_state, strict=True) self.to(self.device).eval().requires_grad_(False) @staticmethod def _validate_format(metadata: dict) -> None: if metadata.get("format") != "yue2sem-tokenizer-1": raise ValueError("Expected a unified yue2sem-tokenizer-1 checkpoint") if metadata.get("head.kind") != "cvq": raise ValueError("This inference package supports CVQ tokenizer checkpoints") @staticmethod def _json_object(metadata: dict, key: str) -> dict: try: value = json.loads(metadata[key]) except (KeyError, json.JSONDecodeError, TypeError) as exc: raise ValueError(f"Checkpoint has invalid or missing {key}") from exc if not isinstance(value, dict): raise ValueError(f"{key} must contain a JSON object") return value @staticmethod def _integer(values: dict, key: str, minimum: int) -> int: try: value = int(values[key]) except (KeyError, TypeError, ValueError) as exc: raise ValueError(f"Checkpoint has invalid or missing {key}") from exc if isinstance(values[key], float) and not values[key].is_integer(): raise ValueError(f"{key} must be an integer") if value < minimum: raise ValueError(f"{key} must be at least {minimum}") return value @staticmethod def _boolean(metadata: dict, key: str) -> bool: value = metadata.get(key) if value not in ("0", "1"): raise ValueError(f"{key} must be '0' or '1'") return value == "1" @torch.inference_mode() def tokenize(self, waveform: torch.Tensor, batch_windows: int = 8) -> np.ndarray: """Convert fp32 [channels, samples] 48 kHz audio to int32 tokens at 25 Hz.""" import torchaudio if waveform.ndim != 2 or waveform.shape[0] < 1: raise ValueError("Expected audio shaped [channels, samples]") if batch_windows < 1: raise ValueError("batch_windows must be positive") mono = torchaudio.functional.resample( waveform.float().mean(0), 48000, ENCODER_RATE ).to(self.device) features = self.encoder.features(mono, self.layer, self.chunk) # The deployed pipeline transfers fp16 features before head inference. features = features.half().cpu().numpy() return self._predict(features, batch_windows).astype(np.int32) def _predict(self, features: np.ndarray, batch_windows: int) -> np.ndarray: x = features.astype(np.float32) if self.instnorm: x = (x - x.mean(0)) / (x.std(0) + 1e-5) x = torch.from_numpy(x) total = len(x) starts = list(range(0, max(1, total - HEAD_WINDOW + 1), HEAD_WINDOW // 2)) if starts[-1] + HEAD_WINDOW < total: starts.append(max(0, total - HEAD_WINDOW)) output = torch.zeros(total, dtype=torch.long) for first in range(0, len(starts), batch_windows): group = starts[first : first + batch_windows] windows = torch.stack( [ F.pad( x[start : start + HEAD_WINDOW], (0, 0, 0, HEAD_WINDOW - min(HEAD_WINDOW, total - start)), ) for start in group ] ).to(self.device) with torch.autocast( self.device.type, dtype=torch.bfloat16, enabled=self.device.type == "cuda", ): logits = self.head(windows).float() for start, window_logits in zip(group, logits): count = min(HEAD_WINDOW, total - start) low = start + (0 if start == 0 else HEAD_WINDOW // 4) high = start + count - (0 if start + count >= total else HEAD_WINDOW // 4) prediction = window_logits[:count].argmax(-1) output[low:high] = prediction[low - start : high - start].cpu() return to_token_clock(output.numpy(), self.offset)